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

    val test : unit -> 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 =
        | DBot
        | DVal of int
        | DChoice of id * (item * edd) list

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

          

    (*
      edd_print : Format.formatter -> edd_v -> unit
    *)
    let edd_print fmt v =
        let ids = Hashtbl.create 12 in
        let n = ref 0 in
        let rec mkid = function
            | DChoice(_, l) as x ->
              if not (Hashtbl.mem ids x) then begin
                incr n; Hashtbl.add ids x !n;
                List.iter (fun (_, x) -> mkid x) l
              end
            | _ -> ()
        in mkid v.root;
        let print_n fmt = function
          | DBot -> Format.fprintf fmt "⊥";
          | DVal i -> Format.fprintf fmt "v%d" i;
          | DChoice _ as x -> Format.fprintf fmt "n%d" (Hashtbl.find ids x)
        in
        Format.fprintf fmt "Root: %a@." print_n v.root;
        Hashtbl.iter
            (fun x id -> 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: %s ? @[<hov>%a@]@." id var
                    (print_list p ", ") l
              | _ -> assert false)
            ids;
        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 = 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 = 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 root = match r with
        | EItem x ->
          DChoice(vid,
            List.map (fun v -> if op v x then v, DVal 0 else v, DBot)
              (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 nb x =
              DChoice(b,
                List.map (fun v -> if op v x then v, DVal 0 else v, DBot)
                  (List.assoc b ve.evars))
            in
            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.copy a.leaves in
        let n = ref 0 in
        Hashtbl.iter (fun i _ -> if i > !n then n := i) leaves;

        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 memo_bcopy = Hashtbl.create 12 in
        let rec bcopy nb =
          try Hashtbl.find memo_bcopy nb
          with Not_found -> let r =
            match nb with
            | DBot -> DBot
            | DVal i -> get_leaf (Hashtbl.find b.leaves i)
            | DChoice (v, x) ->
              DChoice(v, List.map (fun (c, n) -> c, bcopy n) x)
          in Hashtbl.add memo_bcopy nb r; r
        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, x -> bcopy x
            | 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 Hashtbl.add memo (na, nb) r; r
        in
        { leaves; ve; root = 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_of_conslist : varenv -> enum_cons list -> edd_v

      Does this by dichotomy so that the EDD meets are done mostly on
      small sized EDDs.
    *)
    let rec edd_of_conslist ve = function
      | [] -> edd_top ve
      | [a] -> edd_of_cons ve a
      | l ->
        let rec split = function
          | [] -> [], []
          | [a] -> [a], []
          | a::b::q -> let u, v = split q in a::u, b::v
        in
        let u, v = split l in
        edd_meet (edd_of_conslist ve u) (edd_of_conslist ve v)


    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 *)
    }

    (*
      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 (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]
          | _ -> []
        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 lv = lv0 @ extract_linked_evars init_cl @ extract_linked_evars cl_g in
        let lv = uniq_sorted
             (List.sort Pervasives.compare (List.map ord_couple lv)) in
        let evars_ord = scope_constrict evars (lv0 @ lv) 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

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

          


    let do_prog opt rp =
        let e = init_env opt rp in
        ()

end