open Ast
open Ast_util
open Formula
open Typing

open Util
open Num_domain
open Abs_interp

module I (ND : NUMERICAL_ENVIRONMENT_DOMAIN) : sig

    val do_prog : cmdline_opt -> rooted_prog -> unit

end = struct


    (*  **********************
              EDD Domain  
        **********************  *)

    type item = string

    type evar = id * item list
    type nvar = id * bool

    type varenv = {
        evars         : evar list;
        nvars         : nvar list;
        ev_order      : (id, int) Hashtbl.t;
    }

    type edd_node =
        | DBot
        | DVal of int
        | DChoice of id * (item * edd) list
    and edd = int * edd_node

    type edd_v = {
        ve            : varenv;
        root          : edd;
        leaves        : (int, ND.t) Hashtbl.t;
        (* add here eventual annotations *)
    }

    (*
      Utility functions for memoization

      memo : (((int * 'a) -> (int * 'b)) -> 'a -> 'b) -> (int * 'a) -> (int * 'b)
      memo2 : (((int * 'a) -> (int * 'b) -> (int * 'c)) -> 'a -> 'b -> 'c) -> (int * 'a) -> (int * 'b) -> (int * 'c)
    *)
    let memo f =
        let n = ref 0 in
        let memo = Hashtbl.create 12 in
        let rec ff (id, v) =
          try Hashtbl.find memo id
          with Not_found ->
            let r = f ff v in
            incr n;
            Hashtbl.add memo id (n, r);
            n, r
        in ff
          
    let memo2 f =
        let n = ref 0 in
        let memo = Hashtbl.create 12 in
        let rec ff (id1, v1) (id2, v2) =
          try Hashtbl.find memo (id1, id2)
          with Not_found ->
            let r = f ff v1 v2 in
            incr n;
            Hashtbl.add memo (id1, id2) (n, r);
            n, r
        in ff

    (*
      edd_print : Format.formatter -> edd_v -> unit
    *)
    let edd_print fmt v =
        let c_nodes = Hashtbl.create 12 in
        let rec add = function
            | n, (DChoice(_, l) as x) ->
              if not (Hashtbl.mem c_nodes n) then begin
                Hashtbl.add c_nodes n x;
                List.iter (fun (_, y) -> add y) l
              end
            | _ -> ()
        in add v.root;

        let print_n fmt = function
          | (_, DBot) -> Format.fprintf fmt "⊥";
          | (_, DVal i) -> Format.fprintf fmt "v%d" i;
          | (n, DChoice _) -> Format.fprintf fmt "n%d" n
        in
        Format.fprintf fmt "Root: %a@." print_n v.root;

        Hashtbl.iter
          (fun id x -> match x with
              | DChoice (var, l) ->
                let p fmt (c, l) = Format.fprintf fmt "%s → %a" c print_n l in
                Format.fprintf fmt "n%d: %a ? @[<hov>%a@]@."
                    id Formula_printer.print_id var
                    (print_list p ", ") l
              | _ -> assert false)
          c_nodes;
        Hashtbl.iter
            (fun id v -> Format.fprintf fmt "v%d: %a@." id ND.print v)
            v.leaves
            


    (*
      edd_bot : varenv -> edd_v
    *)
    let edd_bot ve = { ve; root = (0, DBot); leaves = Hashtbl.create 1 }

    (*
      edd_top : evar list -> nvar list -> edd_v
    *)
    let edd_top ve =
        let leaves = Hashtbl.create 1 in
        Hashtbl.add leaves 0 (ND.top ve.nvars);
        { ve; root = (1, DVal 0); leaves }

    (*
      edd_of_cons : varenv -> enum_cons -> edd_v
    *)
    let edd_of_cons ve (op, vid, r) =
        let op = match op with | E_EQ -> (=) | E_NE -> (<>) in

        let leaves = Hashtbl.create 1 in
        Hashtbl.add leaves 0 (ND.top ve.nvars);

        let bot, top = (0, DBot), (1, DVal 0) in

        let root = match r with
        | EItem x ->
          2, DChoice(vid,
              List.map (fun v -> if op v x then v, bot else v, top)
                (List.assoc vid ve.evars))      
        | EIdent vid2 ->
            let a, b =
                if Hashtbl.find ve.ev_order vid < Hashtbl.find ve.ev_order vid2
                  then vid, vid2
                  else vid2, vid
            in
            let n = ref 2 in
            let nb x =
              incr n;
              !n, DChoice(b,
                    List.map (fun v -> if op v x then v, top else v, bot)
                      (List.assoc b ve.evars))
            in
            2, DChoice(a, List.map (fun x -> x, nb x) (List.assoc a ve.evars))
        in
        { ve; root; leaves }

    (*
      edd_join : edd_v -> edd_v -> edd_v
    *)
    let edd_join a b =
        let ve = a.ve in

        let leaves = Hashtbl.create 12 in

        let n = ref 0 in
        let get_leaf x = 
          let id = ref None in
          Hashtbl.iter (fun i v -> if ND.eq v x then id := Some i) leaves;
          begin match !id with
          | Some i -> DVal i
          | None ->
            incr n;
            Hashtbl.add leaves !n x;
            DVal (!n)
          end
        in


        let f f_rec na nb =
          let dq v l =
            let _, (i, x0) = List.hd l in
            if List.exists (fun (_, (j, x)) -> j <> i || x <> x0) l
              then DChoice(v, l)
              else x0
          in
          match na, nb with
          | DBot, DBot -> DBot
          | DBot, DVal i -> get_leaf (Hashtbl.find b.leaves i)
          | x, DBot -> x

          | DChoice(va, la), DChoice(vb, lb) when va = vb ->
            let kl = List.map2
                (fun (ta, ba) (tb, bb) -> assert (ta = tb);
                  ta, f (ba, bb))
                la lb
            in
            dq va kl
          | DChoice(va, la), DChoice(vb, lb) ->
            let v, kl =
              if Hashtbl.find ve.ev_order va < Hashtbl.find ve.ev_order vb then
                va, List.map (fun (k, ca) -> k, f (ca, nb)) la
              else
                vb, List.map (fun (k, cb) -> k, f (na, cb)) lb
            in
            dq v kl

          | DChoice(va, la), _ ->
            let kl = List.map (fun (k, ca) -> k, f (ca, nb)) la in
            dq va kl
          | _, DChoice(vb, lb) ->
            let kl = List.map (fun (k, cb) -> k, f (na, cb)) lb in
            dq vb kl

          | DVal u, DVal v ->
            let x = ND.join (Hashtbl.find a.leaves u) (Hashtbl.find b.leaves v) in
            get_leaf x
        in
        { leaves; ve; root = memo2 f a.root b.root }

    (*
      edd_meet : edd_v -> edd_v -> edd_v
    *)
    let edd_meet a b =
        let ve = a.ve in

        let n = ref 0 in
        let leaves = Hashtbl.create 12 in

        let get_leaf x = 
          if ND.is_bot x then DBot else begin
            let id = ref None in
            Hashtbl.iter (fun i v -> if ND.eq v x then id := Some i) leaves;
            match !id with
            | Some i -> DVal i
            | None ->
              incr n;
              Hashtbl.add leaves !n x;
              DVal (!n)
          end
        in

        let memo = Hashtbl.create 12 in
        let rec f (na, nb) =
          try Hashtbl.find memo (na, nb)
          with Not_found -> let r =
            let dq v l =
              let _, x0 = List.hd l in
              if List.exists (fun (_, x) -> x <> x0) l
                then DChoice(v, l)
                else x0
            in
            match na, nb with
            | DBot, _ | _, DBot -> DBot
            | DChoice(va, la), DChoice(vb, lb) when va = vb ->
              let kl = List.map2
                  (fun (ta, ba) (tb, bb) -> assert (ta = tb);
                    ta, f (ba, bb))
                  la lb
              in
              dq va kl
            | DChoice(va, la), DChoice(vb, lb) ->
              let v, kl =
                if Hashtbl.find ve.ev_order va < Hashtbl.find ve.ev_order vb then
                  va, List.map (fun (k, ca) -> k, f (ca, nb)) la
                else
                  vb, List.map (fun (k, cb) -> k, f (na, cb)) lb
              in
              dq v kl
            | DChoice(va, la), _ ->
              let kl = List.map (fun (k, ca) -> k, f (ca, nb)) la in
              dq va kl
            | _, DChoice(vb, lb) ->
              let kl = List.map (fun (k, cb) -> k, f (na, cb)) lb in
              dq vb kl
            | DVal u, DVal v ->
              let x = ND.meet (Hashtbl.find a.leaves u) (Hashtbl.find b.leaves v) in
              get_leaf x
          in Hashtbl.add memo (na, nb) r; r
        in
        { leaves; ve; root = f (a.root, b.root) }

    (*
      edd_apply_ncl : edd_v -> num_cons list -> edd_v
    *)
    let edd_apply_ncl v ncl =
        let ve = v.ve in

        let n = ref 0 in
        let leaves = Hashtbl.create 12 in

        let get_leaf x = 
          if ND.is_bot x then DBot else begin
            let id = ref None in
            Hashtbl.iter (fun i v -> if ND.eq v x then id := Some i) leaves;
            match !id with
            | Some i -> DVal i
            | None ->
              incr n;
              Hashtbl.add leaves !n x;
              DVal (!n)
          end
        in

        let memo = Hashtbl.create 12 in
        let rec f n =
          try Hashtbl.find memo n
          with Not_found -> let r =
            match n with
            | DBot -> DBot
            | DVal i -> get_leaf (ND.apply_cl (Hashtbl.find v.leaves i) ncl)
            | DChoice(var, l) ->
              let l = List.map (fun (k, v) -> k, f v) l in
              let _, x0 = List.hd l in
              if List.exists (fun (_, x) -> x <> x0) l
                then DChoice(var, l)
                else x0
          in Hashtbl.add memo n r; r
        in
        { leaves; ve; root = f v.root }

    (*
      edd_apply_cl : edd_v -> conslist -> edd_v
    *)
    let rec edd_apply_cl v (ec, nc, r) =
        let v = List.fold_left
            (fun v cons -> edd_meet v (edd_of_cons v.ve cons))
            v ec in
        match r with
        | CLTrue ->
          edd_apply_ncl v nc
        | CLFalse -> edd_bot v.ve
        | CLAnd (a, b) ->
          let v = edd_apply_cl v ([], nc, a) in
          edd_apply_cl v ([], nc, b)
        | CLOr((eca, nca, ra), (ecb, ncb, rb)) ->
          edd_join (edd_apply_cl v (eca, nc@nca, ra))
                   (edd_apply_cl v (ecb, nc@ncb, rb))

    (*
      edd_eq : edd_v -> edd_v -> bool
    *)
    let edd_eq a b =
        let memo = Hashtbl.create 12 in
        let rec f (na, nb) =
          try Hashtbl.find memo (na, nb)
          with Not_found -> let r =
            match na, nb with
            | DBot, DBot -> true
            | DVal i, DVal j ->
              ND.eq (Hashtbl.find a.leaves i) (Hashtbl.find b.leaves j)
            | DChoice(va, la), DChoice(vb, lb) when va = vb ->
              List.for_all2 (fun (ca, na) (cb, nb) -> assert (ca = cb); f (na, nb))
                  la lb
            | _ -> false
          in Hashtbl.add memo (na, nb) r; r
        in f (a.root, b.root)


    (*
    let test () =
        let ve = {
            evars = ["x", ["tt"; "ff"]; "y", ["tt"; "ff"]; "z", ["tt"; "ff"]];
            nvars = [];
            ev_order = Hashtbl.create 2 } in
        Hashtbl.add ve.ev_order "x" 0;
        Hashtbl.add ve.ev_order "y" 1;
        Hashtbl.add ve.ev_order "z" 2;
        let u = edd_of_cons ve (E_EQ, "x", EIdent "y") in
        Format.printf "x = y : @[%a@]@." edd_print u;
        let v = edd_of_cons ve (E_NE, "y", EIdent "z") in
        Format.printf "y != z : @[%a@]@." edd_print v;
        let w = edd_meet u v in
        Format.printf "x = y && y != z : @[%a@]@." edd_print w;
        let t = edd_join u v in
        Format.printf "x = y || y != z : @[%a@]@." edd_print t
    *)


    (* ******************************
            Abstract interpret
      ******************************* *)
    
    type env = {
        rp            : rooted_prog;
        opt           : cmdline_opt;

        ve            : varenv;

        (* program expressions *)
        f             : bool_expr;
        cl            : conslist;
        f_g           : bool_expr;
        cl_g          : conslist;
        guarantees    : (id * bool_expr) list;

        (* abstract interpretation *)
        cycle         : (id * id * typ) list;     (* s'(x) = s(y) *)
        forget        : (id * typ) list;          (* s'(x) not specified *)
        mutable data  : edd_v;
    }

    (*
      extract_linked_evars : conslist -> (id * id) list

      Extract all pairs of enum-type variable (x, y) appearing in an
      equation like x = y or x != y

      A couple may appear several times in the result.
    *)
    let rec extract_linked_evars_root (ecl, _, r) =
        let v_ecl = List.fold_left
            (fun c (_, x, v) -> match v with
                | EIdent y -> (x, y)::c
                | _ -> c)
            [] ecl
        in
        v_ecl

    let rec extract_linked_evars(ecl, _, r) =
        let v_ecl = List.fold_left
            (fun c (_, x, v) -> match v with
                | EIdent y -> (x, y)::c
                | _ -> c)
            [] ecl
        in
        let v_ecl2 =
          let q = List.fold_left
            (fun c (_, x, v) -> match v with
              | EItem _ -> x::c | _ -> c)
            [] ecl
          in
          match q with
          | [x; y] -> [x, y]
          | [x; y; z] -> [x, y; y, z; z, x]
          | _ -> []
        in
        let rec aux = function
          | CLTrue | CLFalse -> []
          | CLAnd(a, b) -> aux a @ aux b
          | CLOr(a, b) -> extract_linked_evars a @ extract_linked_evars b
        in
        v_ecl @ v_ecl2 @ aux r

    (*
      scope_constrict : id list -> (id * id) list -> id list

      Orders the variable in the first argument such as to minimize the
      sum of the distance between the position of two variables appearing in
      a couple of the second list. (minimisation is approximate, this is
      an heuristic so that the EDD will not explode in size when expressing
      equations such as x = y && u = v && a != b)
    *)
    let scope_constrict vars cp_id =
        let var_i = Array.of_list vars in
        let n = Array.length var_i in

        let i_var = Hashtbl.create n in
        Array.iteri (fun i v -> Hashtbl.add i_var v i) var_i;

        let cp_i = List.map
          (fun (x, y) -> Hashtbl.find i_var x, Hashtbl.find i_var y)
          cp_id in

        let eval i =
          let r = Array.make n (-1) in
          Array.iteri (fun pos var -> r.(var) <- pos) i;
          Array.iteri (fun _ x -> assert (x <> (-1))) r;
          List.fold_left
            (fun s (x, y) -> s + abs (r.(x) - r.(y)))
            0 cp_i
        in

        let best = Array.make n 0 in
        for i = 0 to n-1 do best.(i) <- i done;

        let usefull = ref true in
        Format.printf "SCA";
        while !usefull do
          Format.printf ".@?";

          usefull := false;
          let try_s x =
            if eval x < eval best then begin
              Array.blit x 0 best 0 n;
              usefull := true
            end
          in

          for i = 0 to n-1 do
            let tt = Array.copy best in
            (* move item i at beginning *)
            let temp = tt.(i) in
            for j = i downto 1 do tt.(j) <- tt.(j-1) done;
            tt.(0) <- temp;
            (* try all positions *)
            try_s tt;
            for j = 1 to n-1 do
              let temp = tt.(j-1) in
              tt.(j-1) <- tt.(j);
              tt.(j) <- temp;
              try_s tt
            done
          done
        done;
        Format.printf "@.";

        Array.to_list (Array.map (Array.get var_i) best)

    (*
      init_env : cmdline_opt -> rooted_prog -> env
    *)
    let init_env opt rp =
        Format.printf "Vars: @[<hov>%a@]@.@."
            (print_list Ast_printer.print_typed_var ", ")
            rp.all_vars;

        let enum_vars = List.fold_left
            (fun v (_, id, t) -> match t with
                | TEnum ch -> (id, ch)::v | _ -> v)
            [] rp.all_vars in
        let num_vars = List.fold_left
            (fun v (_, id, t) -> match t with
                | TInt -> (id, false)::v | TReal -> (id, true)::v | _ -> v)
            [] rp.all_vars in

        let init_f = Transform.init_f_of_prog rp in
        Format.printf "Init formula: %a@.@." Formula_printer.print_expr init_f;
        let init_cl = conslist_of_f init_f in

        let guarantees = Transform.guarantees_of_prog rp in
        Format.printf "Guarantees:@.";
        List.iter (fun (id, f) ->
            Format.printf "  %s: %a@." id Formula_printer.print_expr f)
          guarantees;
        Format.printf "@.";

        let f = Formula.eliminate_not (Transform.f_of_prog rp false) in
        let f_g = Formula.eliminate_not (Transform.f_of_prog rp true) in
        Format.printf "Cycle formula:@.%a@.@." Formula_printer.print_expr f;

        let cl = Formula.conslist_of_f f in
        let cl_g = Formula.conslist_of_f f_g in

        (* calculate order for enumerated variables *)
        let evars = List.map fst enum_vars in

        let lv0 = List.map (fun x -> x, "N"^x)
            (List.filter (fun x -> List.exists (fun y -> y = "N"^x) evars) evars) in
        let lv1 = extract_linked_evars_root init_cl
            @ extract_linked_evars_root cl_g in
        let lv2 = extract_linked_evars init_cl @ extract_linked_evars cl_g in

        let lv1 = uniq_sorted
             (List.sort Pervasives.compare (List.map ord_couple lv1)) in
        let lv2 = uniq_sorted
             (List.sort Pervasives.compare (List.map ord_couple lv2)) in
        let evars_ord = scope_constrict evars (lv0 @ lv1 @ lv2) in

        let ev_order = Hashtbl.create (List.length evars) in
        List.iteri (fun i x -> Hashtbl.add ev_order x i) evars_ord;
        let ve = { evars = enum_vars; nvars = num_vars; ev_order } in

        Format.printf "Order for variables: @[<hov>[%a]@]@."
          (print_list Formula_printer.print_id ", ") evars_ord;

        (* calculate cycle variables and forget variables *)
        let cycle = List.fold_left
          (fun q (_, id, ty) ->
              if id.[0] = 'N' then
                (String.sub id 1 (String.length id - 1), id, ty)::q
              else q)
          [] rp.all_vars
        in
        let forget = List.map (fun (_, id, ty) -> (id, ty))
            (List.filter
              (fun (_, id, _) ->
                not (List.exists (fun (_, id2, _) -> id2 = "N"^id) rp.all_vars))
              rp.all_vars)
        in

        (* calculate initial environment *)
        let data = edd_apply_cl (edd_top ve) init_cl in
        Format.printf "Init: @[<hov>%a@]@." edd_print data;

        { rp; opt; ve;
          f; cl; f_g; cl_g; guarantees;
          cycle; forget; data }

          


    let do_prog opt rp =
        let e = init_env opt rp in
        let d2 = edd_apply_cl e.data e.cl in
        Format.printf " -> @[<hov>%a@]@." edd_print d2

end