(***********************************************************************)
(*                                                                     *)
(*                               TAToo                                 *)
(*                                                                     *)
(*                     Kim Nguyen, LRI UMR8623                         *)
(*                   Université Paris-Sud & CNRS                       *)
(*                                                                     *)
(*  Copyright 2010-2013 Université Paris-Sud and Centre National de la *)
(*  Recherche Scientifique. All rights reserved.  This file is         *)
(*  distributed under the terms of the GNU Lesser General Public       *)
(*  License, with the special exception on linking described in file   *)
(*  ../LICENSE.                                                        *)
(*                                                                     *)
(***********************************************************************)

INCLUDE "utils.ml"
INCLUDE "debug.ml"

open Format
open Misc

type stats = { mutable pass : int;
               tree_size : int;
               mutable fetch_trans_cache_access : int;
               mutable fetch_trans_cache_miss : int;
               mutable eval_trans_cache_access : int;
               mutable eval_trans_cache_miss : int;
               mutable nodes_per_run : int list;
             }


let dummy_set = StateSet.singleton State.dummy_state



  IFDEF HTMLTRACE
  THEN
type sat_array = StateSet.t array list
  DEFINE IFHTML(a,b) = (a)
  ELSE
type sat_array = StateSet.t array
  DEFINE IFHTML(a,b) = (b)
  END

let unsafe_get a i =
  if i < 0 then StateSet.empty else
    Array.unsafe_get (IFHTML(List.hd a, a)) i

let unsafe_set a i v  =
  (* if v != old_v then *)
    Array.unsafe_set (IFHTML(List.hd a, a)) i v

type 'node td_action = 'node -> int -> 'node -> 'node -> StateSet.t
type 'node bu_action = 'node -> int -> StateSet.t
let dummy_action = fun _ _ -> assert false

type ('tree, 'node) run = {
  tree : 'tree ;
     (* The argument of the run *)
  auto : Ata.t;
     (* The automaton to be run *)
  mutable sat: sat_array;
     (* A mapping from node preorders to states satisfied at that node *)
  mutable pass : int;
     (* Number of run we have performed *)
  mutable fetch_trans_cache : Ata.Formula.t Cache.N2.t;
     (* A cache from states * label to list of transitions *)
  mutable td_cache : 'node td_action Cache.N6.t;
  mutable bu_cache : 'node bu_action Cache.N6.t;
     (* Two 6-way caches used during the top-down and bottom-up phase
        label * self-set * fc-set * ns-set * parent-set * node-shape -> self-set
     *)
  stats : stats;
}

let dummy_form = Ata.Formula.stay State.dummy_state

let get_form run tag (q : State.t) =
  let auto = run.auto in
  let fetch_trans_cache = run.fetch_trans_cache in
  let stats = run.stats in
  let phi =
    stats.fetch_trans_cache_access <- stats.fetch_trans_cache_access + 1;
    Cache.N2.find fetch_trans_cache (tag.QName.id :> int) (q :> int)
  in
  if phi == dummy_form then
    let phi = Ata.get_form auto tag q in
    let () =
      stats.fetch_trans_cache_miss <- stats.fetch_trans_cache_miss + 1;
      Cache.N2.add
        fetch_trans_cache
        (tag.QName.id :> int)
        (q :> int) phi
    in phi
  else
    phi



let eval_form phi fcs nss ps ss summary =
  let open Ata in
  let rec loop phi =
    begin match Formula.expr phi with
    | Boolean.False -> false
    | Boolean.True -> true
    | Boolean.Atom (a, b) ->
      begin
        let open Tree.NodeSummary in
        match a.Atom.node with
        | Move (m, q) ->
          b && StateSet.mem q (
            match m with
              `First_child -> fcs
            | `Next_sibling -> nss
            | `Parent | `Previous_sibling -> ps
            | `Stay -> ss
          )
        | Is_first_child -> b == is_left summary
        | Is_next_sibling -> b == is_right summary
        | Is k -> b == (k == kind summary)
        | Has_first_child -> b == has_left summary
        | Has_next_sibling -> b == has_right summary
      end
    | Boolean.And(phi1, phi2) -> loop phi1 && loop phi2
    | Boolean.Or (phi1, phi2) -> loop phi1 || loop phi2
    end
  in
  loop phi


let eval_trans_aux run tag summary fcs nss ps sat todo  =
  StateSet.fold (fun q (a_sat) ->
    let phi =
      get_form run tag q
    in
    if eval_form phi fcs nss ps a_sat summary then
      StateSet.add q a_sat
    else a_sat
  ) todo sat


let rec eval_trans_fix run tag summary fcs nss ps sat todo  =
  let new_sat =
    eval_trans_aux run tag summary fcs nss ps sat todo
  in
  if new_sat == sat then sat else
    eval_trans_fix run tag summary fcs nss ps new_sat todo


let eval_trans run trans_cache tag summary fcs nss ps ss todo action_builder =
  let stats = run.stats in
  let fcsid = (fcs.StateSet.id :> int) in
  let nssid = (nss.StateSet.id :> int) in
  let psid = (ps.StateSet.id :> int) in
  let ssid = (ss.StateSet.id :> int) in
  let tagid = (tag.QName.id :> int) in

  let res = Cache.N6.find trans_cache tagid summary ssid fcsid nssid psid in
  stats.eval_trans_cache_access <- 1 + stats.eval_trans_cache_access;
  if res != dummy_action then
    res
  else
    let new_sat =
      eval_trans_fix run tag summary fcs nss ps ss todo
    in
    stats.eval_trans_cache_miss <- 1 + stats.eval_trans_cache_miss;
    let new_action = action_builder ps summary tag new_sat in
    Cache.N6.add trans_cache tagid summary ssid fcsid nssid psid new_action;
    new_action

module NodeSummary = Tree.NodeSummary
module Make (T : Tree.S)  =
struct
  module Tree : Tree.S with type node = T.node = T
  module ResultSet : Deque.S with type elem = Tree.node =
    Deque.Make (struct type t = Tree.node end)


  let make auto tree =
    let len = Tree.size tree in
    {
      tree = tree;
      auto = auto;
      sat = (let a = Array.make len StateSet.empty in
             IFHTML([a], a));
      pass = 0;
      fetch_trans_cache = Cache.N2.create dummy_form;
      td_cache = Cache.N6.create dummy_action;
      bu_cache = Cache.N6.create dummy_action;
      stats = {
        pass = 0;
        tree_size = len;
        fetch_trans_cache_access = 0;
        fetch_trans_cache_miss = 0;
        eval_trans_cache_access = 0;
        eval_trans_cache_miss = 0;
        nodes_per_run = [];
      }
    }


  let top_down run mk_update_res =
    let num_visited = ref 0 in
    let i = run.pass in
    let tree = run.tree in
    let auto = run.auto in
    let states_by_rank = Ata.get_states_by_rank auto in
    let td_todo = states_by_rank.(i) in
    let bu_todo =
      if i == Array.length states_by_rank - 1 then StateSet.empty
      else
        states_by_rank.(i+1)
    in
    let run_sat = run.sat in
    let last_run = i >= Array.length states_by_rank - 2 in

    let rec common node parent parent_sat action_builder =
      begin
        incr num_visited;
        let tag = Tree.tag tree node in
        let node_id = Tree.preorder tree node in
        let summary = Tree.summary tree node in
        let fc = Tree.first_child tree node in
        let ns = Tree.next_sibling tree node in
        (* We enter the node from its parent *)
        let status0 = unsafe_get run_sat node_id in
        (* get the node_statuses for the first child, next sibling and parent *)
        (* evaluate the transitions with all this statuses *)
        let action  =
          eval_trans run
            run.td_cache tag
            summary
            (unsafe_get run_sat (Tree.preorder tree fc))
            (unsafe_get run_sat (Tree.preorder tree ns))
            parent_sat
            status0 td_todo action_builder
        in
        action node node_id fc ns
      end

    and td_action_builder parent_sat summary tag status1 =
      let update_res = mk_update_res false status1 in
      match NodeSummary.(has_left summary, has_right summary) with
        false, false ->
        (fun node node_id fc ns ->
           unsafe_set run_sat node_id status1;
           update_res node;
           StateSet.empty)
      | true, false ->
        (fun node node_id fc ns ->
           unsafe_set run_sat node_id status1;
           update_res node;
           loop_td fc node status1)
      | false, true ->
        (fun node node_id fc ns ->
           unsafe_set run_sat node_id status1;
           update_res node;
           loop_td ns node status1)
      | _ ->
        (fun node node_id fc ns ->
           unsafe_set run_sat node_id status1; (* write the td_states *)
           update_res node;
           ignore (loop_td fc node status1);
           loop_td ns node status1 (* tail call *)
        )
    and td_and_bu_action_builder parent_sat summary tag status1  =
      match NodeSummary.(has_left summary, has_right summary) with
        false, false ->
        (fun node node_id fc ns ->
           let action =
             eval_trans run run.bu_cache tag summary StateSet.empty StateSet.empty
               parent_sat status1 bu_todo bu_action_builder
           in
           action node node_id)
      | true, false ->
        (fun node node_id fc ns ->
           let fcs1 = loop_td_and_bu fc node status1 in
           let action =
             eval_trans run run.bu_cache tag summary fcs1 StateSet.empty
               parent_sat status1 bu_todo bu_action_builder
           in
           action node node_id
        )
      | false, true ->
        (fun node node_id fc ns ->
           let nss1 = loop_td_and_bu ns node status1 in
           let action =
             eval_trans run run.bu_cache tag summary StateSet.empty nss1
               parent_sat status1 bu_todo bu_action_builder
           in
           action node node_id
        )
      | _ ->
        (fun node node_id fc ns ->
           let fcs1 = loop_td_and_bu fc node status1 in
           let nss1 = loop_td_and_bu ns node status1 in
           let action =
             eval_trans run run.bu_cache tag summary fcs1 nss1 parent_sat status1 bu_todo bu_action_builder
           in
           action node node_id
        )
    and bu_action_builder parent_sat summary tag status2 node node_id =
      unsafe_set run_sat node_id status2;
      status2
    and td_and_bu_last_action_builder parent_sat summary tag status1 =
      match NodeSummary.(has_left summary, has_right summary) with
        false, false ->
        (fun node node_id fc ns ->
           let action =
             eval_trans run run.bu_cache tag summary StateSet.empty StateSet.empty
               parent_sat status1 bu_todo bu_last_action_builder
           in
           action node node_id)
      | true, false ->
        (fun node node_id fc ns ->
           let fcs1 = loop_td_and_bu_last fc node status1 in
           let action =
             eval_trans run run.bu_cache tag summary fcs1 StateSet.empty
               parent_sat status1 bu_todo bu_last_action_builder
           in
           action node node_id
        )
      | false, true ->
        (fun node node_id fc ns ->
           let nss1 = loop_td_and_bu_last ns node status1 in
           let action =
             eval_trans run run.bu_cache tag summary StateSet.empty nss1
               parent_sat status1 bu_todo bu_last_action_builder
           in
           action node node_id
        )
      | _ ->
        (fun node node_id fc ns ->
           let nss1 = loop_td_and_bu_last ns node status1 in
           let fcs1 = loop_td_and_bu_last fc node status1 in
           let action =
             eval_trans run run.bu_cache tag summary fcs1 nss1 parent_sat status1 bu_todo bu_last_action_builder
           in
           action node node_id
        )
    and bu_last_action_builder parent_sat summary tag status2 =
      let update_res = mk_update_res true status2 in
      (fun node node_id ->
         unsafe_set run_sat node_id status2;
         update_res node;
         status2)
    and loop_td node parent parent_sat =
      common node parent parent_sat td_action_builder
    and loop_td_and_bu node parent parent_sat =
      common node parent parent_sat td_and_bu_action_builder
    and loop_td_and_bu_last node parent parent_sat =
      common node parent parent_sat td_and_bu_last_action_builder
    in
    let _ =
      if bu_todo == StateSet.empty then
        loop_td (Tree.root tree) Tree.nil dummy_set
      else if last_run then
        loop_td_and_bu_last (Tree.root tree) Tree.nil dummy_set
      else
        loop_td_and_bu (Tree.root tree) Tree.nil dummy_set
    in
    run.pass <- run.pass + 2;
    run.stats.pass <- run.stats.pass + 1;
    run.stats.nodes_per_run <- !num_visited :: run.stats.nodes_per_run



  let mk_update_result auto =
    let sel_states = Ata.get_selecting_states auto in
    let res = ResultSet.create () in
    (fun prepend sat ->
       if StateSet.intersect sat sel_states then
         if prepend then (fun n -> ResultSet.push_front n res)
         else (fun n -> ResultSet.push_back n res)
       else (fun _ -> ())),
    (fun () -> res)


  let mk_update_full_result auto =
    let sel_states = Ata.get_selecting_states auto in
    let res_mapper =
      StateSet.fold_right (fun q acc -> (q, ResultSet.create())::acc) sel_states []
    in
    (fun prepend sat ->
       StateSet.fold (fun q f_acc ->
           if StateSet.mem q sel_states then
             let res = List.assoc q res_mapper in
             if prepend then (fun n -> ResultSet.push_front n res; f_acc n)
             else (fun n -> ResultSet.push_back n res; f_acc n)
           else f_acc) sat (fun _ -> ())),
    (fun () -> res_mapper)

  let prepare_run run list =
    let tree = run.tree in
    let auto = run.auto in
    let sat = IFHTML((List.hd run.sat), run.sat) in
    let sat0 = Ata.get_starting_states auto in
   ResultSet.iter (fun node ->
      let node_id = Tree.preorder tree node in
      sat.(node_id) <- sat0) list


  let compute_run auto tree nodes mk_update_res =
    let run = make auto tree in
    prepare_run run nodes;
    let rank = Ata.get_max_rank auto in
    while run.pass <= rank do
      top_down run mk_update_res;
      IFHTML((run.sat <- (Array.copy (List.hd run.sat)) :: run.sat), ());
      run.td_cache <- Cache.N6.create dummy_action;
      run.bu_cache <- Cache.N6.create dummy_action;
    done;
    IFHTML((run.sat <- List.tl run.sat), ());
    IFHTML(Html_trace.gen_trace auto run.sat (module T : Tree.S with type t = Tree.t) tree ,());
    run


  let last_stats = ref None

  let full_eval auto tree nodes =
    let mk_update_full,get_full = mk_update_full_result auto in
    let run = compute_run auto tree nodes mk_update_full in
    last_stats := Some run.stats;
    get_full ()

  let eval auto tree nodes =
    let mk_update_res,get_res = mk_update_result auto in
    let run = compute_run auto tree nodes mk_update_res in
    last_stats := Some run.stats;
    get_res ()

  let stats () = match !last_stats with
      Some s -> s.nodes_per_run <- List.rev s.nodes_per_run;s
    | None -> failwith "Missing stats"

end