(***********************************************************************)
(*                                                                     *)
(*                               TAToo                                 *)
(*                                                                     *)
(*                   Lucca Hirschi, LRI UMR8623                        *)
(*                   Université Paris-Sud & CNRS                       *)
(*                                                                     *)
(*  Copyright 2010-2012 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.                                                        *)
(*                                                                     *)
(***********************************************************************)

module Node  =
struct
  type t = int
  let hash n = n
  let compare = (-)
  let equal = (=)
end
  
module NodeHash = Hashtbl.Make (Node)
  
type t = (StateSet.t*StateSet.t) NodeHash.t
(** Map from nodes to query and recognizing states *)
(* Note that we do not consider the nil nodes *)

exception Oracle_fail
exception Over_max_fail
exception Max_fail

(* Mapped sets for leaves *)
let map_leaf asta = (Asta.bot_states asta, StateSet.empty)
let empty = (StateSet.empty,StateSet.empty)

(* Build the Oracle *)
let rec bu_oracle asta run tree tnode =
  let node = Tree.preorder tree tnode in
  if Tree.is_leaf tree tnode
  then
    if tnode == Tree.nil
    then ()
    else NodeHash.add run node (map_leaf asta)
  else
    let tfnode = Tree.first_child tree tnode (* first child *)
    and tnnode = Tree.next_sibling tree tnode in (* next-sibling *)
    let fnode,nnode =
      (Tree.preorder tree tfnode, Tree.preorder tree tnnode) in
    begin
      bu_oracle asta run tree tfnode;
      bu_oracle asta run tree tnnode;
      let q_rec n =
        try NodeHash.find run n
        with Not_found -> map_leaf asta in
      let (_,qfr),(_,qnr) = q_rec fnode,q_rec nnode (* computed in rec call *)
      and lab = Tree.tag tree tnode in
      let _,list_tr = Asta.transitions_lab asta lab in (* only take reco. *)
      let rec result set = function
        | [] -> set
        | (q,form) :: tl ->
          if Formula.eval_form (qfr,qnr) form
          then result (StateSet.add q set) tl
          else result set tl in
      let result_set = result StateSet.empty list_tr in
      NodeHash.add run node (StateSet.empty, result_set)
    end

(* Build the over-approx. of the maximal run *)
let rec bu_over_max asta run tree tnode =
  if (Tree.is_leaf tree tnode)      (* BU_oracle has already created the map *)
  then
    ()
  else
    let tfnode = Tree.first_child tree tnode
    and tnnode = Tree.next_sibling tree tnode in
    begin
      bu_over_max asta run tree tfnode;
      bu_over_max asta run tree tnnode;
      let (fnode,nnode) =
        (Tree.preorder tree tfnode, Tree.preorder tree tnnode)
      and node = Tree.preorder tree tnode in          
      let q_rec n =
        try NodeHash.find run n
        with Not_found -> map_leaf asta in
      let (qfq,qfr),(qnq,qnr) = q_rec fnode,q_rec nnode in
      let lab = Tree.tag tree tnode in
      let list_tr,_ = Asta.transitions_lab asta lab in (* only take query st. *)
      let rec result set = function
        | [] -> set
        | (q,form) :: tl ->
          if Formula.infer_form (qfq,qnq) (qfr,qnr) form
          then result (StateSet.add q set) tl
          else result set tl in
      let _,resultr = try NodeHash.find run node
        with _ -> raise Over_max_fail in      
      let result_set = result StateSet.empty list_tr in
      NodeHash.replace run node (result_set, resultr)
    (* Never remove elt in Hash (the old one would appear) *)
    end


(* Build the maximal run *)
let rec tp_max asta run tree tnode =
  if (Tree.is_leaf tree tnode)      (* BU_oracle has already created the map *)
  then
    ()
  else
    let node = Tree.preorder tree tnode
    and tfnode = Tree.first_child tree tnode
    and tnnode = Tree.next_sibling tree tnode in
    let (fnode,nnode) =
      (Tree.preorder tree tfnode, Tree.preorder tree tnnode) in
    let q_rec n =
      try NodeHash.find run n
      with Not_found -> (Asta.bot_states asta,StateSet.empty) in
    let (qf),(qn) = q_rec fnode,q_rec nnode in
    let lab = Tree.tag tree tnode in
    let list_tr,_ = Asta.transitions_lab asta lab in (* only take query. *)
    let set_node,_ = try NodeHash.find run node
      with _ -> raise Max_fail in
    let rec result = function
      | [] -> []
      | (q,form) :: tl ->
        if (Formula.infer_form qf qn form) && (StateSet.mem q set_node)
        then form :: (result tl)
        else result tl in
    let list_form = result list_tr in
    let rec add_st (ql,qr) = function
      | [] -> ql,qr
      | f :: tl -> let sql,sqr = Formula.st f in
                   let ql' = StateSet.union sql ql
                   and qr' = StateSet.union sqr qr in
                   add_st (ql',qr') tl in
    let ql,qr = add_st (StateSet.empty, StateSet.empty) list_form in
    let qfq,qfr = try NodeHash.find run fnode
      with | _ -> map_leaf asta
    and qnq,qnr = try NodeHash.find run nnode
      with | _ -> map_leaf asta in
    begin
      if Tree.is_leaf tree tfnode
      then ()
      else NodeHash.replace run fnode (StateSet.inter qfq ql,qfr);
      if Tree.is_leaf tree tnnode
      then ()
      else NodeHash.replace run nnode (StateSet.inter qnq qr,qnr);
      tp_max asta run tree tfnode;
      tp_max asta run tree tnnode;
    end

let compute tree asta =
  let flag = 2 in                       (* debug  *)
  let size_tree = 10000 in              (* todo *)
  let map = NodeHash.create size_tree in
  bu_oracle asta map tree (Tree.root tree);
  if flag > 0 then begin
    bu_over_max asta map tree (Tree.root tree);
    if flag = 2
    then
      tp_max asta map tree (Tree.root tree)
    else ()
  end
  else ();
  map

let selected_nodes tree asta =
  let run = compute tree asta in
  NodeHash.fold
    (fun key set acc ->
      if not(StateSet.is_empty
               (StateSet.inter (fst set) (Asta.selec_states asta)))        
      then key :: acc
      else acc)
    run []

let print fmt run =
  let print_d_set fmt (s_1,s_2) =
    Format.fprintf fmt "(%a,%a)"
      StateSet.print s_1 StateSet.print s_2 in
  let print_map fmt run =
    let pp = Format.fprintf fmt in
    if NodeHash.length run = 0
    then Format.fprintf fmt "ø"
    else
      NodeHash.iter (fun cle set -> pp "|  %i->%a  @ " cle print_d_set set)
        run in
  let print_box fmt run =
    let pp = Format.fprintf fmt in
    pp "@[<hov 0>@.  # Mapping:@.   @[<hov 0>%a@]@]"
      print_map run
  in
  Format.fprintf fmt "@[<hov 0>##### RUN #####@, %a@]@." print_box run