open Ast
open Ast_util
open Formula

open Util
open Environment_domain

module I (E : ENVIRONMENT_DOMAIN) : sig

    type st

    val do_prog     : int -> prog -> id -> unit

end = struct

    type st = {
        p           : prog;
        widen_delay : int;
        root_scope  : scope;
        all_vars    : var_def list;
        env         : E.t;
        f           : bool_expr;
        cl          : conslist;
        f_g         : bool_expr;
        cl_g        : conslist;
        guarantees  : (id * bool_expr) list;
    }


    let node_vars p f nid =
      let (n, _) = find_node_decl p f in
      List.map (fun (p, id, t) -> p, nid^"/"^id, t)
        (n.args @ n.ret @ n.var)

    (*
        extract_all_vars : prog -> scope -> var_decl list

        Extracts all variables with names given according to
        naming convention used here and in transform.ml :
        - pre variables are not prefixed by scope
        - next values for variables are prefixed by capital N
    *)
    let rec extract_all_vars p (node, prefix, eqs) =
        let vars_of_expr e =
          List.flatten
            (List.map
              (fun (f, id, eqs, args) ->
                node_vars p f (node^"/"^id) @
                extract_all_vars p (node^"/"^id, "", eqs))
              (extract_instances p e))
          @
          List.flatten
            (List.map
              (fun (id, _) -> [false, id, AST_TINT; false, "N"^id, AST_TINT])
                  (* TODO : type of pre value ? *)
              (extract_pre e))
        in
        let vars_of_eq e = match fst e with
          | AST_assign(_, e) | AST_assume(_, e) | AST_guarantee(_, e) -> vars_of_expr e
          | AST_activate (b, _) ->
            let rec do_branch = function
              | AST_activate_body b ->
                List.map (fun (p, id, t) -> p, node^"/"^id, t) b.act_locals
                @
                extract_all_vars p (node, b.act_id^".", b.body)
              | AST_activate_if(c, a, b) ->
                vars_of_expr c @ do_branch a @ do_branch b
            in do_branch b
          | AST_automaton _ -> not_implemented "extract_all vars automaton"
        in
        (false, node^"/"^prefix^"time", AST_TINT)::
        (false, "N"^node^"/"^prefix^"time", AST_TINT)::
        (List.flatten (List.map vars_of_eq eqs))

    (* init state *)
    let init_state widen_delay p root =
      let root_scope = get_root_scope p root in

      let f = Formula.eliminate_not (Transform.f_of_prog p root false) in
      let cl = Formula.conslist_of_f f in
      let f_g = Formula.eliminate_not (Transform.f_of_prog p root true) in
      let cl_g = Formula.conslist_of_f f_g in

      let all_vars = node_vars p root "" @ extract_all_vars p root_scope in
      let env = List.fold_left
        (fun env (_, id, ty) ->
          E.addvar env id (ty = AST_TREAL))
        E.init
        all_vars
      in
      let init_f = Transform.init_f_of_prog p root in
      let guarantees = Transform.guarantees_of_prog p root in

      Format.printf "Init formula: %a@.@." Formula_printer.print_expr init_f;
      Format.printf "Cycle formula:@.%a@.@." Formula_printer.print_expr f;
      Format.printf "Cycle formula with guarantees:@.%a@.@." Formula_printer.print_expr f_g;

      Format.printf "Vars: @[<hov>%a@]@.@." (Ast_printer.print_list Ast_printer.print_var_decl ", ") all_vars;

      Format.printf "Guarantees:@.";
      List.iter (fun (id, f) ->
          Format.printf "  %s: %a@." id Formula_printer.print_expr f)
        guarantees;
      Format.printf "@.";

      let env = E.apply_f env init_f in

      { p; root_scope; widen_delay;
        f; cl; f_g; cl_g;
        guarantees; all_vars; 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), NIdent 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.all_vars (E.apply_cl i cl)) 

    (*
      loop : st -> env -> env -> env
    *)
    let loop st pnew stay =

      Format.printf "Loop @[<hv>%a@]@.New @[<hv>%a@]@.Stay @[<hv>%a@]@."
        Formula_printer.print_conslist st.cl
        E.print_all pnew E.print_all stay;

      let add_stay = cycle st st.cl pnew in
      let acc0 = E.join stay add_stay in

      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
      y, z

    (*
      do_prog : prog -> id -> unit
    *)
    let do_prog widen_delay p root =
      let st = init_state widen_delay p root in

      let pnew = st.env in

      let stay, stay_c = loop st pnew (E.vbottom pnew) in

      let final = E.join (E.apply_cl pnew st.cl) stay_c in

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

      let check_guarantee (id, f) =
        Format.printf "@[<hv>%s:@ %a@ "
          id Formula_printer.print_expr f;
        let z = E.apply_f final (BAnd(f, st.f)) in
        if E.is_bot z then
          Format.printf "OK@]@ "
        else
          Format.printf "FAIL@]@ "
      in
      Format.printf "Guarantees: @[<v 0>";
      List.iter check_guarantee st.guarantees;
      Format.printf "@]@.";

      Format.printf "Probes:@[<v>";
      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.all_vars;
      Format.printf "@]@."

end