open Ast
open Ast_util
open Formula
open Typing
open Cmdline

open Util
open Num_domain
open Enum_domain

open Varenv

module I (ED : ENUM_ENVIRONMENT_DOMAIN) (ND : NUMERICAL_ENVIRONMENT_DOMAIN)
: sig

    val do_prog : cmdline_opt -> rooted_prog -> unit

end = struct

    type abs_v = ED.t * ND.t

    type location = {
        id            : int;

        def           : bool_expr list;   (* conjunction of formula *)
        def_cl        : conslist;
        is_init       : bool;

        f             : bool_expr;
        cl            : conslist;

        (* For chaotic iteration fixpoint *)
        mutable star  : bool;
        mutable in_c  : int;
        mutable v     : abs_v;

        mutable out_t : int list;
        mutable in_t  : int list;
    }

    type env = {
        rp            : rooted_prog;
        opt           : cmdline_opt;

        ve            : varenv;

        (* program expressions *)
        f             : bool_expr;
        guarantees    : (id * bool_expr) list;

        (* data *)
        loc           : (int, location) Hashtbl.t;
        counter       : int ref;
    }

    (* **************************
            ABSTRACT VALUES
       ************************** *)

    (*
      top : env -> abs_v
      bottom : env -> abs_v
    *)
    let top e = (ED.top e.ve.evars, ND.top e.ve.nvars)
    let bottom e = (ED.top e.ve.evars, ND.bottom e.ve.nvars)
    let is_bot (e, n) = ND.is_bot n

    let print_v fmt (enum, num) =
      if is_bot (enum, num) then
        Format.fprintf fmt "⊥"
      else
        Format.fprintf fmt "@[<hov 4>(%a,@ %a)@]" ED.print enum ND.print num

    (*
      join : abs_v -> abs_v -> abs_v
      meet : abs_v -> abs_v -> abs_v
    *)
    let join (e1, n1) (e2, n2) =
      if is_bot (e1, n1) then (e2, n2)
      else if is_bot (e2, n2) then (e1, n1)
      else (ED.join e1 e2, ND.join n1 n2)
    let meet (e1, n1) (e2, n2) =
      if is_bot (e1, n1) then ED.vtop e1, ND.vbottom n1
      else if is_bot (e2, n2) then ED.vtop e2, ND.vbottom n2
      else
        try (ED.meet e1 e2 , ND.meet n1 n2)
        with Bot -> e1, ND.vbottom n1

    (*
      apply_cl : abs_v -> conslist -> abs_v
    *)
    let rec apply_cl (enum, num) (ec, nc, r) =
        try
          begin match r with
          | CLTrue ->
            (ED.apply_cl enum ec, ND.apply_cl num nc)
          | CLFalse ->
            (enum, ND.vbottom num)
          | CLAnd(a, b) ->
            let enum, num = apply_cl (enum, num) (ec, nc, a) in
            let enum, num = apply_cl (enum, num) ([], nc, b) in
            enum, num
          | CLOr((eca, nca, ra), (ecb, ncb, rb)) ->
            let a = apply_cl (enum, num) (ec@eca, nc@nca, ra) in 
            let b = apply_cl (enum, num) (ec@ecb, nc@ncb, rb) in 
            join a b
          end
        with Bot -> ED.vtop enum, ND.vbottom num


    (* ***************************
              INTERPRET
       *************************** *)



    (*
      init_env : cmdline_opt -> rooted_prog -> env
    *)
    let init_env opt rp =
        let f = Transform.f_of_prog_incl_init rp false in

        let f = simplify_k (get_root_true f) f in
        Format.printf "Complete formula:@.%a@.@." Formula_printer.print_expr f;

        let facts = get_root_true f in
        let f, rp, _ = List.fold_left
          (fun (f, (rp : rooted_prog), repls) eq ->
            match eq with
            | BEnumCons(E_EQ, a, EIdent b)
                when a.[0] <> 'L' && b.[0] <> 'L' ->

              let a = try List.assoc a repls with Not_found -> a in
              let b = try List.assoc b repls with Not_found -> b in

              if a = b then
                f, rp, repls
              else begin
                let keep, repl =
                  if String.length a <= String.length b
                    then a, b
                    else b, a
                in
                Format.printf "Replacing %s with %s@." repl keep;
                let f = formula_replace_evars [repl, keep; "L"^repl, "L"^keep] f in
                let rp =
                  { rp with all_vars =
                      List.filter (fun (_, id, _) -> id <> repl) rp.all_vars } in
                let repls = (repl, keep)::
                  (List.map (fun (r, k) -> r, if k = repl then keep else k) repls) in
                f, rp, repls
              end
            | _ -> f, rp, repls)
          (f, rp, []) facts in

        let f = simplify_k (get_root_true f) f in

        Format.printf "Complete formula after simpl:@.%a@.@."
            Formula_printer.print_expr f;

        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 ve = mk_varenv rp f (conslist_of_f f) in

        let env = {
          rp; opt; ve; f; guarantees;
          loc = Hashtbl.create 2; counter = ref 2; } in

        (* add initial disjunction : L/must_reset = tt, L/must_reset ≠ tt *)
        let rstc = [BEnumCons(E_EQ, "L/must_reset", EItem bool_true)] in
        let rstf = simplify_k rstc f in
        let rstf = simplify_k (get_root_true rstf) rstf in
        let nrstc = [BEnumCons(E_NE, "L/must_reset", EItem bool_true)] in
        let nrstf = simplify_k nrstc f in
        let nrstf = simplify_k (get_root_true nrstf) nrstf in
        Hashtbl.add env.loc 0
          {
            id = 0;
            def = rstc;
            def_cl = conslist_of_f (f_and_list rstc);
            is_init = true;

            f = rstf;
            cl = conslist_of_f rstf;

            star = false;
            in_c = 0;
            v = (ED.top ve.evars, ND.bottom ve.nvars);

            out_t = [];
            in_t = [];
          };
        Hashtbl.add env.loc 1
          {
            id = 1;
            def = nrstc;
            def_cl = conslist_of_f (f_and_list nrstc);
            is_init = false;

            f = nrstf;
            cl = conslist_of_f nrstf;

            star = false;
            in_c = 0;
            v = (ED.top ve.evars, ND.bottom ve.nvars);

            out_t = [];
            in_t = [];
          };
        env


    (*
      pass_cycle : env -> edd_v -> edd_v
      unpass_cycle : env -> edd_v -> edd_v

      set_target_case : env -> edd_v -> bool_expr -> edd_v
      cycle : env -> edd_v -> conslist -> edd_v
    *)
    let pass_cycle env (enum, num) =
        let assign_e, assign_n = List.fold_left
          (fun (ae, an) (a, b, t) -> match t with
            | TEnum _ -> (a, b)::ae, an
            | TInt | TReal -> ae, (a, NIdent b)::an)
          ([], []) env.cycle in

        let enum = ED.assign enum assign_e in
        let num = ND.assign num assign_n in

        let ef, nf = List.fold_left
            (fun (ef, nf) (var, t) -> match t with
              | TEnum _ -> var::ef, nf
              | TReal | TInt -> ef, var::nf)
            ([], []) env.forget in

        (ED.forgetvars enum ef, List.fold_left ND.forgetvar num nf)

    let unpass_cycle env (enum, num) =
        let assign_e, assign_n = List.fold_left
          (fun (ae, an) (a, b, t) -> match t with
            | TEnum _ -> (b, a)::ae, an
            | TInt | TReal -> ae, (b, NIdent a)::an)
          ([], []) env.ve.cycle in

        let enum = ED.assign enum assign_e in
        let num = ND.assign num assign_n in

        let ef, nf = List.fold_left
            (fun (ef, nf) (var, t) -> match t with
              | TEnum _ -> var::ef, nf
              | TReal | TInt -> ef, var::nf)
            ([], []) env.ve.forget_inv in

        (ED.forgetvars enum ef, List.fold_left ND.forgetvar num nf)

    let set_target_case e v cl =
        let tgt = (unpass_cycle e (apply_cl (top e) cl)) in
        (*Format.printf "Target %a = %a@." Formula_printer.print_conslist cl print_v tgt;*)
        meet v tgt

    (*
      chaotic_iter : env -> unit

      Fills the values of loc[*].v, and updates out_t and in_t
    *)
    let chaotic_iter e =
        let delta = ref [] in

        (* Fill up initial states *)
        Hashtbl.iter
          (fun q loc ->
            loc.out_t <- [];
            loc.in_t <- [];
            if loc.is_init then begin
              loc.v <- apply_cl (top e) loc.cl;
              delta := q::!delta
            end else
              loc.v <- bottom e)
          e.loc;

        (* Iterate *)
        let it_counter = ref 0 in
        while !delta <> [] do
          let s = List.hd !delta in
          let loc = Hashtbl.find e.loc s in

          incr it_counter;
          Format.printf "@.Iteration %d: q%d@." !it_counter s;

          let start = loc.v in
          let f i =
            (*Format.printf "I: %a@." print_v i;*)
            let i' = set_target_case e i loc.def_cl in
            (*Format.printf "I': %a@." print_v i';*)
            let j = join start (apply_cl (apply_cl (pass_cycle e.ve i') loc.def_cl) loc.cl) in
            (*Format.printf "J: %a@." print_v j;*)
            j
          in

          let rec iter n (i_enum, i_num) =
            let fi_enum, fi_num = f (i_enum, i_num) in
            let j_enum, j_num =
              if ND.is_bot fi_num then
                i_enum, i_num
              else
                if n < e.opt.widen_delay then
                  ED.join i_enum fi_enum, ND.join i_num fi_num
                else
                  ED.join i_enum fi_enum, ND.widen i_num fi_num
            in
            if ED.eq j_enum i_enum && ND.eq j_num i_num
              then (i_enum, i_num)
              else iter (n+1) (j_enum, j_num)
          in
          let y = iter 0 start in
          let z = fix (fun (a, b) (c, d) -> ED.eq a c && ND.eq b d) f y in
          Format.printf "Fixpoint: %a@." print_v z;

          loc.v <- z;

          Hashtbl.iter
            (fun t loc2 ->
                let u = pass_cycle e.ve z in
                let v = apply_cl u loc2.def_cl in
                let w = apply_cl v loc2.cl in
                (*Format.printf "u: %a@.v: %a@. w: %a@." print_v u print_v v print_v w;*)
                let r_enum, r_num = w in
                if not (is_bot (r_enum, r_num)) then begin
                  Format.printf "%d -> %d with:@.  %a@." s t print_v (r_enum, r_num);
                  if not (List.mem s loc2.in_t)
                    then loc2.in_t <- s::loc2.in_t;
                  if not (List.mem t loc.out_t)
                    then loc.out_t <- t::loc.out_t;
                  let enum, num = loc2.v in
                  let enum2, num2 = join (enum, num) (r_enum, r_num) in
                  if not (ED.subset enum2 enum) || not (ND.subset num2 num) then
                  begin
                    loc2.v <- (enum2, num2);
                    if not (List.mem t !delta)
                      then delta := t::!delta
                  end
                end)
            e.loc;

          delta := List.filter ((<>) s) !delta;
        done


    let print_locs e =
        Hashtbl.iter
          (fun id loc ->
            Format.printf "@.";
            Format.printf "q%d: @[<hov 4>[ %a ]@]@." id
              (print_list Formula_printer.print_expr " ∧ ") loc.def;
            (*Format.printf "F: %a@." Formula_printer.print_conslist loc.cl;*)
            Format.printf "    @[<hv 0>%a ∧@ %a@]@."
              ED.print (fst loc.v) ND.print (snd loc.v);
            Format.printf " -> @[<hov>[%a]@]@."
              (print_list (fun fmt i -> Format.fprintf fmt "q%d" i) ", ") loc.out_t;
            )
          e.loc

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

        Format.printf "@.Initializing.@.";
        chaotic_iter e;
        print_locs e

end