open Ast
open Ast_util
open Formula

open Util
open Environment_domain

module I (E : ENVIRONMENT_DOMAIN) : sig

    type st

    val do_prog     : prog -> id -> unit

end = struct

    type st = {
        p           : prog;
        root_scope  : scope;
        all_vars    : var_def list;
        env         : E.t;
        f           : bool_expr;
        cl          : conslist;
    }


    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^"/"^b.act_id^"."^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 p root =
      let root_scope = get_root_scope p root in

      let f = Formula.eliminate_not (Transform.f_of_prog p root) in
      let cl = Formula.conslist_of_f f 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 env = E.apply_f env init_f in

      Format.printf "Init formula: %a@.@." Formula_printer.print_expr init_f;
      Format.printf "Cycle formula: %a@.@." Formula_printer.print_conslist cl;

      { p; root_scope; f; cl; 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
      Format.printf "Pass cycle: %a@.@." E.print_all e;
      e


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

      let rec step acc i n =
        Format.printf "Step %d: %a@." n E.print_all acc;
        if n < 10 then begin
          let i' = E.apply_cl (E.join i acc) st.cl in
          let j = pass_cycle st.all_vars i' in
          let acc' = (if n > 6 then E.widen else E.join) acc j in
          step acc' i (n+1)
        end
      in
      step (E.vbottom st.env) st.env 0
end