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;
        depth           : int;

        mutable def     : abs_v;
        is_init         : bool;

        mutable f       : bool_expr;
        mutable cl      : conslist;

        (* For chaotic iteration fixpoint *)
        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) = ED.is_bot e || ND.is_bot n

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

    (*
      join : abs_v -> abs_v -> abs_v
      widen : abs_v -> abs_v -> abs_v
      meet : abs_v -> abs_v -> abs_v
    *)
    let join a b =
      if is_bot a then b
      else if is_bot b then a
      else (ED.join (fst a) (fst b), ND.join (snd a) (snd b))

    let widen a b =
      if is_bot a then b
      else if is_bot b then a
      else (ED.join (fst a) (fst b), ND.widen (snd a) (snd b))

    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 -> ED.vtop e1, ND.vbottom n1

    (*
      eq_v : abs_v -> abs_v -> bool
      subset_v : abs_v -> abs_v -> bool
    *)
    let eq_v (a, b) (c, d) =
        (is_bot (a, b) && is_bot (c, d)) 
          || (ED.eq a c && ND.eq b d)

    let subset_v (a, b) (c, d) =
      (is_bot (a, b)) ||
        (not (is_bot (c, d)) && ED.subset a c && ND.subset b d)

    (*
      apply_cl : abs_v -> conslist -> abs_v
    *)
    let rec apply_cl (enum, num) (ec, nc, r) =
        begin match r with
        | CLTrue ->
          begin
            try (ED.apply_cl enum ec, ND.apply_cl num nc)
            with Bot -> ED.vtop enum, ND.vbottom num
          end
        | CLFalse ->
          (ED.vtop 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

    (*
      apply_cl_all_cases : abs_v -> conslist -> abs_v list
    *)
    let rec apply_cl_all_cases v (ec, nc, r) =
        match r with
        | CLTrue ->
          let v =
            try ED.apply_cl (fst v) ec, ND.apply_cl (snd v) nc
            with Bot -> ED.vtop (fst v), ND.vbottom (snd v)
          in
          if is_bot v then [] else [v]
        | CLFalse ->
          []
        | CLAnd(a, b) ->
          let q1 = apply_cl_all_cases v (ec, nc, a) in
          List.flatten
            (List.map (fun c -> apply_cl_all_cases c ([], [], b)) q1)
        | CLOr((eca, nca, ra), (ecb, ncb, rb)) ->
          let la = apply_cl_all_cases v (ec@eca, nc@nca, ra) in 
          let lb = apply_cl_all_cases v (ec@ecb, nc@ncb, rb) in 
          lb@(List.filter (fun a -> not (List.exists (fun b -> eq_v a b) lb)) la)


    (* ***************************
              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;
            depth = 0;
            def = apply_cl (top env) (conslist_of_f rstc);
            is_init = true;

            f = rstf;
            cl = conslist_of_f rstf;

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

            out_t = [];
            in_t = [];
          };
        Hashtbl.add env.loc 1
          {
            id = 1;
            depth = 0;
            def = apply_cl (top env) (conslist_of_f nrstc);
            is_init = false;

            f = nrstf;
            cl = conslist_of_f nrstf;

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

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


    (*
      ternary_conds : bool_expr -> bool_expr list
    *)
    let rec ternary_conds = function
        | BAnd(a, b) -> ternary_conds a @ ternary_conds b
        | BTernary(c, a, b) -> [c]
        | _ -> []

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

    (*
      print_locs : env -> unit
    *)

    let print_locs_defs e =
        Hashtbl.iter
          (fun id loc ->
            Format.printf "q%d: @[<v 2>%a@]@." id print_v loc.def;
            )
          e.loc

    let print_locs e =
        Hashtbl.iter
          (fun id loc ->
            Format.printf "@.";
            Format.printf "q%d (depth = %d):@.  D: @[<v 2>%a@]@." id loc.depth print_v loc.def;
            (*Format.printf "  F: (%a)@." Formula_printer.print_expr loc.f;*)
            Format.printf "  V: %a@." print_v loc.v;
            Format.printf " -> @[<hov>[%a]@]@."
              (print_list (fun fmt i -> Format.fprintf fmt "q%d" i) ", ") loc.out_t;
            )
          e.loc

    (*
      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 <- [];
            loc.in_c <- 0;
            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;

        print_locs_defs e;

        (* 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' = meet i (unpass_cycle e loc.def) in
            (*Format.printf "I': %a@." print_v i';*)
            let j = join start
              (apply_cl
                (meet (pass_cycle e.ve i') loc.def)
                loc.cl) in
            (*Format.printf "J: %a@." print_v j;*)
            j
          in

          let rec iter n i =
            let fi = f i in
            let j = 
                if n < e.opt.widen_delay then
                  join i fi
                else
                  widen i fi
            in
            if eq_v i j
              then i
              else iter (n+1) j
          in
          let y = iter 0 start in
          let z = fix eq_v f y in
          let u = pass_cycle e.ve z in

          if e.opt.verbose_ci then
            Format.printf "Fixpoint: %a@. mem fp: %a@." print_v z print_v u;

          loc.v <- z;

          Hashtbl.iter
            (fun t loc2 ->
                let v = meet u loc2.def 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;*)
                if not (is_bot w) then begin
                  if e.opt.verbose_ci then
                    Format.printf "%d -> %d with:@.  %a@." s t print_v w;
                  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;
                  if not (subset_v w loc2.v) then begin
                    if loc2.in_c < e.opt.widen_delay then
                      loc2.v <- join loc2.v w
                    else
                      loc2.v <- widen loc2.v w;
                    loc2.in_c <- loc2.in_c + 1;
                    if not (List.mem t !delta)
                      then delta := t::!delta
                  end
                end)
            e.loc;

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

        (* remove useless locations *)
        let useless = ref [] in
        Hashtbl.iter
          (fun i loc ->
            if is_bot loc.v then begin
              Format.printf "Useless location detected: q%d@." i;
              useless := i::!useless
            end)
          e.loc;
        List.iter (Hashtbl.remove e.loc) !useless;

        print_locs e



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

        let rec iter n =
          Format.printf "@.--------------@.Refinement #%d@." n;

          chaotic_iter e;

          let qc = ref None in

          (* put true or false conditions into location definition *)
          Hashtbl.iter
            (fun q (loc : location) ->
              let rec iter () =
                try
                  let cond = List.find
                    (fun c ->
                      is_bot (apply_cl loc.v (conslist_of_f c))
                      || is_bot (apply_cl loc.v (conslist_of_f (BNot c))))
                    (ternary_conds loc.f)
                  in
                  let tr =
                    if is_bot (apply_cl loc.v (conslist_of_f cond))
                    then BNot cond
                    else cond
                  in
                  loc.def <- apply_cl loc.def (conslist_of_f tr);
                  loc.f <- simplify_k [tr] loc.f;
                  loc.f <- simplify_k (get_root_true loc.f) loc.f;
                  loc.cl <- conslist_of_f loc.f;
                  iter()
                with Not_found -> ()
              in iter ())
            e.loc;

          (* find splitting condition *)
          Hashtbl.iter
            (fun q (loc:location) ->
              if !qc = None && loc.depth < e.opt.max_dp_depth then
              let cs = ternary_conds loc.f in
              List.iter
                (fun c ->
                  let cases_t = apply_cl_all_cases (top e) (conslist_of_f c) in
                  let cases_f = apply_cl_all_cases (top e) (conslist_of_f (BNot c)) in
                  let cases = cases_t @ cases_f in
                  if
                    List.length 
                      (List.filter
                        (fun case -> not (is_bot (meet loc.v case)))
                        cases)
                      >= 2
                  &&
                    (List.exists
                        (fun qi ->
                          let loci = Hashtbl.find e.loc qi in
                          let v = apply_cl
                            (meet (pass_cycle e.ve loci.v) loc.def)
                            loc.cl in
                          List.exists
                            (fun case -> is_bot (meet v case))
                            cases)
                        loc.in_t
                    || List.exists
                        (fun case ->
                          let v = meet loc.v case in
                          List.exists
                            (fun qo ->
                              let loco = Hashtbl.find e.loc qo in
                              let w = apply_cl
                                (meet (pass_cycle e.ve v) loco.def)
                                loco.cl in
                              is_bot w)
                            loc.out_t)
                        cases)
                  then
                    qc := Some(q, c, cases_t, cases_f)
                  )
                cs
            )
            e.loc;

          if Hashtbl.length e.loc < e.opt.max_dp_width then
            match !qc with
            | None ->
              Format.printf "@.Found no more possible refinement.@.@."
            | Some (q, c, cases_t, cases_f) ->
              Format.printf "@.Refine q%d : @[<v 2>[ %a ]@]@." q
                (print_list print_v ", ") (cases_t@cases_f);

              let loc = Hashtbl.find e.loc q in
              Hashtbl.remove e.loc loc.id;

              let handle_case cc case =
                  if not (is_bot (meet loc.v case)) then
                    let ff = simplify_k [cc] loc.f in
                    let ff = simplify_k (get_root_true ff) ff in
                  
                    let loc2 =
                      { loc with
                        id = (incr e.counter; !(e.counter));
                        depth = loc.depth + 1;
                        def = meet loc.def case;
                        f = ff;
                        cl = conslist_of_f ff } in
                    Hashtbl.add e.loc loc2.id loc2
              in
                List.iter (handle_case c) cases_t;
                List.iter (handle_case (BNot c)) cases_f;
                          
              iter (n+1)
        in iter 0;

        (* Check guarantees *)
        let check_guarantee (id, f) =
          Format.printf "@[<v 4>";
          let cl = Formula.conslist_of_f f in
          Format.printf "%s:@ %a ⇒ ⊥  @ "
            id Formula_printer.print_conslist cl;
          let violate = ref [] in
          Hashtbl.iter
            (fun id loc ->
              if not (is_bot (apply_cl loc.v cl)) then
                violate := id::!violate)
            e.loc;
          if !violate = [] then
            Format.printf "OK"
          else
            Format.printf "VIOLATED in @[<hov 2>[ %a ]@]"
              (print_list (fun fmt i -> Format.fprintf fmt "q%d" i) ", ") !violate;
          Format.printf "@]@ ";
        in
        if e.guarantees <> [] then begin
          Format.printf "Guarantee @[<v 0>";
          List.iter check_guarantee e.guarantees;
          Format.printf "@]@."
        end




end