open Ast
open Ast_util
open Formula
open Typing

open Util
open Abs_domain
open Num_domain

module I (E : ENVIRONMENT_DOMAIN) : sig

    type st

    val do_prog     : int -> rooted_prog -> unit

end = struct

    type st = {
        rp          : rooted_prog;
        widen_delay : int;
        env         : E.t;
        f           : bool_expr;
        cl          : conslist;
        f_g         : bool_expr;
        cl_g        : conslist;
        guarantees  : (id * bool_expr) list;
    }


    (* init state *)
    let init_state widen_delay rp =
      let env = List.fold_left
        (fun env (_, id, ty) -> E.addvar env id ty)
        E.init
        rp.all_vars
      in
      Format.printf "Vars: @[<hov>%a@]@.@."
          (Ast_printer.print_list Ast_printer.print_typed_var ", ")
          rp.all_vars;

      let init_f = Transform.init_f_of_prog rp in
      Format.printf "Init formula: %a@.@." Formula_printer.print_expr init_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 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;
      Format.printf "Cycle formula with guarantees:@.%a@.@." Formula_printer.print_expr f_g;

      let env = E.set_disjunct env "/exit" DNever in
      let env = E.apply_f env init_f in

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

      { rp; widen_delay;
        f; cl; f_g; cl_g;
        guarantees; env }

    (*
      pass_cycle : var_def list -> E.t -> E.t

      Does :
        - x <- Nx, for all x
        - ignore x for all x such that Nx does not exist
    *)
    let pass_cycle vars e =
      let tr_assigns = List.fold_left
        (fun q (_, id, _) ->
            if id.[0] = 'N' then
              (String.sub id 1 (String.length id - 1), id)::q
            else
              q)
        [] vars
      in
      let e = E.assign e tr_assigns in

      let forget_vars = List.map (fun (_, id, _) -> id)
          (List.filter
              (fun (_, id, _) ->
                 not (List.exists (fun (_, id2, _) -> id2 = "N"^id) vars))
              vars) in
      let e = List.fold_left E.forgetvar e forget_vars in
      e


    (* cycle : st -> cl -> env -> env *)
    let cycle st cl i =
      (pass_cycle st.rp.all_vars (E.apply_cl i cl)) 

    (*
      loop : st -> env -> env -> env
    *)
    let loop st acc0 =

      Format.printf "Loop @[<hv>%a@]@.@."
        Formula_printer.print_conslist st.cl;

      let fsharp cl i = 
        Format.printf "Acc @[<hv>%a@]@." E.print_all i;
        let j = cycle st cl i in
        E.join acc0 j
      in

      let rec iter n i =
        let i' =
          if n < st.widen_delay
            then E.join i (fsharp st.cl_g i)
            else E.widen i (fsharp st.cl_g i)
        in
        if E.eq i i' then i else iter (n+1) i'
      in

      let x = iter 0 acc0 in
      let y = fix E.eq (fsharp st.cl_g) x in    (* decreasing iteration *)
      let z = E.apply_cl y st.cl in             (* no more guarantee *)
      y, z

    (*
      do_prog : rooted_prog -> unit
    *)
    let do_prog widen_delay rp =
      let st = init_state widen_delay rp in

      let final_nc, final = loop st st.env in

      Format.printf "@[<hov>F %a@]@.@."
        E.print_all final_nc;
      Format.printf "@[<hov>F' %a@]@.@."
        E.print_all final;

      let check_guarantee (id, f) =
        let cl = Formula.conslist_of_f f in
        Format.printf "@[<hv 4>%s:@ %a ⇒ ⊥  @ "
          id Formula_printer.print_conslist cl;
        let z = E.apply_cl final cl in
        if E.is_bot z then
          Format.printf "OK@]@ "
        else
          Format.printf "FAIL@]@ "
      in
      if st.guarantees <> [] then begin
        Format.printf "Guarantee @[<v 0>";
        List.iter check_guarantee st.guarantees;
        Format.printf "@]@."
      end;

      if List.exists (fun (p, _, _) -> p) st.rp.all_vars then begin
        Format.printf "Probes: @[<v 0>";
        List.iter (fun (p, id, _) ->
          if p then Format.printf "%a ∊ %a@ "
                Formula_printer.print_id id
                E.print_itv (E.var_itv final id))
          st.rp.all_vars;
        Format.printf "@]@."
      end

end