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;

        mutable def     : bool_expr list;   (* conjunction of formula *)
        mutable def_cl  : conslist;
        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) = 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 (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 widen (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.widen 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

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

    let subset_v (a, b) (c, d) =
      ND.is_bot b ||
        (not (ND.is_bot 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) =
        try
          begin match r with
          | CLTrue ->
            let enum =
              fix ED.eq (fun v -> ED.apply_cl v ec) enum
            in
            (enum, 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;

            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;

            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)

    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 <- [];
            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;

        (* 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 =
            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
          (*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;
                  if not (subset_v (r_enum, r_num) loc2.v) then begin
                    if loc2.in_c < e.opt.widen_delay then
                      loc2.v <- join (r_enum, r_num) loc2.v
                    else
                      loc2.v <- widen (r_enum, r_num) loc2.v;
                    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


    let print_locs e =
        Hashtbl.iter
          (fun id loc ->
            Format.printf "@.";
            Format.printf "q%d: @[<v 2>[ %a ]@]@." id
              (print_list Formula_printer.print_expr " ∧ ") loc.def;
            (*Format.printf "  F: (%a)@." Formula_printer.print_expr loc.f;*)
            Format.printf "    %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

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

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

          chaotic_iter e;
          print_locs 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 <- tr::loc.def;
                  loc.def_cl <- conslist_of_f (f_and_list loc.def);
                  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 then
              let cs = ternary_conds loc.f in
              List.iter
                (fun c ->
                  let split_e_case_fold_aux cases c =
                        match c with
                        | BEnumCons(_, x, EItem _) ->
                          (List.map (fun v -> BEnumCons(E_EQ, x, EItem v))
                            (List.assoc x e.ve.evars))@cases
                        | _ -> c::cases
                  in
                  let cases_t =
                    List.fold_left split_e_case_fold_aux []
                      (split_cases [c]) in
                  let cases_f =
                    List.fold_left split_e_case_fold_aux []
                      (split_cases [BNot c]) in
                  let cases = cases_t @ cases_f in
                  if
                    List.length 
                      (List.filter
                        (fun case ->
                          not (is_bot (apply_cl loc.v (conslist_of_f case))))
                        cases)
                      >= 2
                  &&
                    (List.exists
                        (fun qi ->
                          let loci = Hashtbl.find e.loc qi in
                          let v = apply_cl
                            (apply_cl (pass_cycle e.ve loci.v) loc.def_cl)
                            loc.cl in
                          List.exists
                            (fun case -> is_bot (apply_cl v (conslist_of_f case)))
                            cases)
                        loc.in_t
                    || List.exists
                        (fun case ->
                          let v = apply_cl loc.v (conslist_of_f case) in
                          List.exists
                            (fun qo ->
                              let loco = Hashtbl.find e.loc qo in
                              let w = apply_cl
                                (apply_cl (pass_cycle e.ve v) loco.def_cl)
                                loco.cl in
                              is_bot w)
                            loc.out_t)
                        cases)
                  then
                    qc := Some(q, c, cases_t, cases_f)
                  )
                cs
            )
            e.loc;

          (*if n < 7 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 Formula_printer.print_expr ", ") (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 (apply_cl loc.v (conslist_of_f 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));
                        def = case::loc.def;
                        def_cl = conslist_of_f (f_and_list (case::loc.def));
                        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

end