open Ast
open Util
open Ast_util

module I : sig

  exception Bad_datatype
  type value

  val int_value : int -> value
  val bool_value : bool -> value
  val real_value : float -> value

  val as_int : value -> int
  val as_bool : value -> bool
  val as_real : value -> float

  val str_repr_of_val : value -> string

  type state

  val print_state : Format.formatter -> state -> unit

  type io = (id * value) list

  (*
    Get the constants only
  *)
  val consts : prog -> id -> value VarMap.t

  (*
    Construct initial state for a program.
    The id is the root node of the program evaluation.
  *)
  val init_state : prog -> id -> state

  (*
    Run a step of the program (not necessary to specify the program,
      it should be encoded in the state).
    State -> Inputs -> Next state, Outputs
  *)
  val step : state -> io -> (state * io)

end = struct

  (* Values for variables *)

  exception Bad_datatype

  type value =
    | VInt of int
    | VBool of bool
    | VReal of float
    | VState of id
    | VPrevious of value list
    | VBusy   (* intermediate value : calculating ! *)
    | VCalc of (unit -> unit)

  let int_value i = VInt i
  let bool_value b = VBool b
  let real_value r = VReal r

  let as_int = function
    | VInt i -> i
    | _ -> raise Bad_datatype
  let as_bool = function
    | VBool b -> b
    | _ -> raise Bad_datatype
  let as_real = function
    | VReal r -> r
    | _ -> raise Bad_datatype

  let rec str_repr_of_val = function
    | VInt i -> string_of_int i
    | VReal r -> string_of_float r
    | VBool b -> if b then "true" else "false"
    | VState s -> "state " ^ s
    | VPrevious p ->
      "[" ^
        List.fold_left (fun s v ->
            (if s = "" then "" else s ^ ", ") ^ str_repr_of_val v)
          "" p
      ^ "]"
    | VBusy -> "#"
    | VCalc _ -> "()"

  (* States of the interpret *)


  type state = {
    p          : prog;
    root_scope : scope;
    outputs    : id list;
    save       : value VarMap.t;
  }

  let print_state fmt st =
    VarMap.iter
      (fun id v -> Format.fprintf fmt "%s = %s@." id (str_repr_of_val v))
      st.save

  type io = (id * value) list


  (* Internal type env : contains the environment for state calculation *)
  type env = {
    st      : state;
    vars    : (id, value) Hashtbl.t;
  }


  (*
    get_var : env -> id -> id -> value

    Gets the value of a variable relative to a node path.
  *)
  let rec get_var env node id =
    let p = node^"/"^id in
    try match Hashtbl.find env.vars p with
      | VCalc f ->
        Hashtbl.replace env.vars p VBusy;
        f ();
        get_var env node id
      | VBusy -> combinatorial_cycle p
      | x -> x
    with
      | Not_found -> no_variable p

  (*
    eval_expr : env -> (id * id) -> expr ext -> value list
  *)
  let rec eval_expr env (node, prefix) exp =
    let sub_eval = eval_expr env (node, prefix) in
    match fst exp with
    | AST_identifier (id, _) ->
      [loc_error (snd exp) (get_var env node) id]
    | AST_idconst (id, _) ->
      begin try [VarMap.find ("cst/"^id) env.st.save]
      with Not_found -> loc_error (snd exp) no_variable id end
    (* on numerical values *)
    | AST_int_const (i, _) -> [VInt (int_of_string i)]
    | AST_real_const (r, _) -> [VReal (float_of_string r)]
    | AST_unary(AST_UPLUS, e) -> sub_eval e
    | AST_unary(AST_UMINUS, e) ->
      begin match sub_eval e with
      | [VInt x] -> [VInt (-x)]
      | [VReal x] -> [VReal(-. x)]
      | _ -> type_error "Unary on non-numerical."
      end
    | AST_binary(op, e1, e2) ->
      let iop, fop = match op with
      | AST_PLUS -> (+), (+.)
      | AST_MINUS -> (-), (-.)
      | AST_MUL -> ( * ), ( *. )
      | AST_DIV -> (/), (/.)
      | AST_MOD -> (mod), mod_float
      in
      begin match sub_eval e1, sub_eval e2 with
      | [VInt a], [VInt b] -> [VInt (iop a b)]
      | [VReal a], [VReal b] -> [VReal(fop a b)]
      | _ -> type_error "Invalid arguments for numerical binary."
      end
    (* on boolean values *)
    | AST_bool_const b -> [VBool b]
    | AST_binary_rel(op, e1, e2) ->
      let r = match op with
      | AST_EQ -> (=)     | AST_NE -> (<>)
      | AST_LT -> (<)     | AST_LE -> (<=)
      | AST_GT -> (>)     | AST_GE -> (>=)
      in
      begin match sub_eval e1, sub_eval e2 with
      | [VInt a], [VInt b] -> [VBool (r a b)]
      | [VReal a], [VReal b] -> [VBool (r a b)]
      | [VBool a], [VBool b] -> [VBool (r a b)]
      | _ -> type_error "Invalid arguments for binary relation."
      end
    | AST_binary_bool(op, e1, e2) ->
      let r = match op with
      | AST_AND -> (&&)   | AST_OR -> (||)
      in
      begin match sub_eval e1, sub_eval e2 with
      | [VBool a], [VBool b] -> [VBool (r a b)]
      | _ -> type_error "Invalid arguments for boolean relation."
      end
    | AST_not(e) ->
      begin match sub_eval e with
      | [VBool b] -> [VBool (not b)]
      | _ -> type_error "Invalid arguments for boolean negation."
      end
    (* temporal primitives *)
    | AST_pre(exp, n) ->
      begin try match VarMap.find (node^"/"^prefix^n) env.st.save with
        | VPrevious x -> x
        | _ -> assert false
      with Not_found -> []
      end
    | AST_arrow(before, after) -> 
      begin try match VarMap.find (node^"/"^prefix^"init") env.st.save with
      | VBool true -> sub_eval before
      | VBool false -> sub_eval after
      | _ -> assert false
      with Not_found -> assert false
      end
    (* other *)
    | AST_if(cond, a, b) ->
      begin match sub_eval cond with
      | [VBool true] -> sub_eval a
      | [VBool false] -> sub_eval b
      | _ -> type_error "Invalid condition in if."
      end
    | AST_instance((f, _), args, nid) ->
      let (n, _) = find_node_decl env.st.p f in
      List.map
        (fun (_, (id, _), _) -> get_var env (node^"/"^nid) id)
        n.ret

  (*
    reset_scope : env -> scope -> unit

    Sets all the init variables to true, and all the state variables
    to initial state. Does this recursively.
    The state reset done by this on the environment is only applied on the
    next iteration, when the state info has been extracted by extract_st.
  *)
  let rec reset_scope env (node, prefix, eqs) =
    Hashtbl.replace env.vars (node^"/"^prefix^"init") (VBool true);
    let do_exp e =
      List.iter
        (fun (id, eqs, _) -> reset_scope env (node^"/"^id, "", eqs))
        (extract_instances env.st.p e)
    in
    let do_eq (e, _) = match e with
      | AST_assign(_, e) | AST_assume(_, e) | AST_guarantee(_, e) -> do_exp e
      | AST_automaton (aid, states, _) ->
        let (init_state, _) = List.find (fun (st, _) -> st.initial) states in
        Hashtbl.replace env.vars (node^"/"^aid^".state") (VState init_state.st_name);
        List.iter
          (fun (st, _) -> reset_scope env (node, aid^"."^st.st_name^".", st.body))
          states
      | AST_activate (x, _) ->
        let rec aux = function
          | AST_activate_if (_, a, b) -> aux a; aux b
          | AST_activate_body b ->
            reset_scope env (node, b.act_id^".", b.body)
        in aux x
    in
    List.iter do_eq eqs

  (*
    activate_scope : env -> scope -> unit

    Adds functions for calculating the variables in this scope.
    For assign : simple !
    For automaton : lazily select+activate automaton state (would enable
        implementation of strong transitions)
    For activate : lazily select+activate active block
  *)
  let rec activate_scope env (node, prefix, eqs) =
    Hashtbl.replace env.vars (node^"/"^prefix^"act") (VBool true);
    let do_expr e =
      List.iter
        (fun (id, eqs, args) ->
          activate_scope env (node^"/"^id, "", eqs);
          let do_arg ((_,(name,_),_), expr) =
            let apath = node^"/"^id^"/"^name in
            let calc () =
              match eval_expr env (node, prefix) expr with
              | [v] -> Hashtbl.replace env.vars apath v
              | _ -> loc_error (snd expr) error "Unsupported arity for argument."
            in
            Hashtbl.replace env.vars apath (VCalc calc)
          in
          List.iter do_arg args)
        (extract_instances env.st.p e)
    in
    let do_eq (e, _) = match e with
      | AST_assign(vars, e) ->
        do_expr e;
        let calc () =
          let vals = eval_expr env (node, prefix) e in
          List.iter2 (fun (id, _) v ->
            Hashtbl.replace env.vars (node^"/"^id) v) vars vals
        in
        List.iter
          (fun (id, _) ->
              Hashtbl.replace env.vars (node^"/"^id) (VCalc (calc)))
          vars
      | AST_assume(_,e) | AST_guarantee(_,e) -> do_expr e
      | AST_activate(t, vars) ->
        let rec do_tree = function
          | AST_activate_body _ -> ()
          | AST_activate_if(c, a, b) ->
            do_expr c; do_tree a; do_tree b
        in do_tree t;
        let rec needed x () = match x with
          | AST_activate_body b ->
            activate_scope env (node, b.act_id^".", b.body)
          | AST_activate_if (c, t, f) ->
            begin match eval_expr env (node, prefix) c with
            | [VBool true] -> needed t ()
            | [VBool false] -> needed f ()
            | _ -> error "Invalid value in activate if condition."
            end
        in
        List.iter (fun id ->
            Hashtbl.replace env.vars (node^"/"^id) (VCalc (needed t)))
          vars
      | AST_automaton(aid, states, vars) ->
        let asn = match VarMap.find (node^"/"^aid^".state") env.st.save with
          | VState s -> s
          | _ -> assert false
        in
        let (st, _) = List.find (fun (st, _) -> st.st_name = asn) states in
        activate_scope env (node, aid^"."^st.st_name^".", st.body);
        List.iter (fun (c, _, _) -> do_expr c) st.until
    in
    List.iter do_eq eqs

  (*
    is_scope_active : env -> scope -> bool
  *)
  let is_scope_active env (node, prefix, _) =
    try as_bool (Hashtbl.find env.vars (node^"/"^prefix^"act"))
    with Not_found -> false

  (*
    do_weak_transitions : env -> scope -> unit
  *)
  let rec do_weak_transitions env (node, prefix, eqs) =
    let do_expr e =
      List.iter
        (fun (id, eqs, args) ->
          do_weak_transitions env (node^"/"^id, "", eqs))
        (extract_instances env.st.p e)
    in
    let do_eq (e, _) = match e with
      | AST_assign(_, e) -> do_expr e
      | AST_assume (_, e) | AST_guarantee (_, e) -> do_expr e
      | AST_activate(t, _) ->
        let rec do_tree = function
          | AST_activate_body b -> do_weak_transitions env (node, b.act_id^".", b.body)
          | AST_activate_if(c, a, b) ->
            do_expr c; do_tree a; do_tree b
        in do_tree t
      | AST_automaton(aid, states, vars) ->
        let svn = node^"/"^aid^".state" in
        let stv = VarMap.find svn env.st.save in
        let new_st =
          if is_scope_active env (node, prefix, eqs) then begin
            let asn = match stv with
              | VState s -> s
              | _ -> assert false
            in
            let (st, _) = List.find (fun (st, _) -> st.st_name = asn) states in
            try
              let (_, (next_st,_), rst) = List.find
                (fun (c, _, _) -> eval_expr env (node, aid^"."^st.st_name^".") c = [VBool true])
                st.until
              in
                if rst then begin
                  let (st, _) = List.find (fun (st, _) -> st.st_name = next_st) states in
                  reset_scope env (node, aid^"."^next_st^".", st.body)
                end;
                VState(next_st)
            with
              | _ -> stv
          end else
            stv
        in
            Hashtbl.replace env.vars svn new_st
    in
    List.iter do_eq eqs

  (*
    extract_st : env -> state
  *)
  let extract_st env =
    let rec aux (node, prefix, eqs) save =

      let init =
        try as_bool (Hashtbl.find env.vars (node^"/"^prefix^"init"))
        with Not_found ->
          let pre_init = as_bool (VarMap.find (node^"/"^prefix^"init") env.st.save)
          in pre_init && (not (is_scope_active env (node, prefix, eqs)))
      in
      let save = VarMap.add (node^"/"^prefix^"init") (VBool init) save in

      let save_expr save e =
        (* Save pre expressions *)
        let save = List.fold_left
          (fun save (id, expr) ->
            let n = node^"/"^prefix^id in
            if is_scope_active env (node, prefix, eqs) then
              VarMap.add
                n
                (VPrevious (eval_expr env (node, prefix) expr))
                save
            else
              try VarMap.add n
                (VarMap.find n env.st.save)
                save
              with Not_found -> save
          )
          save (extract_pre e) in
        (* Save recursively in sub instances of nodes *)
        let save = List.fold_left
          (fun save (n, eqs, _) ->
            aux (node^"/"^n, "", eqs) save)
          save (extract_instances env.st.p e)
        in save
      in
      let save_eq save eq = match fst eq with
        | AST_assign(_, e) | AST_assume(_, e) | AST_guarantee(_, e) ->
            save_expr save e
        | AST_automaton(aid, states, _) ->
          let svn = node^"/"^aid^".state" in
          let st = Hashtbl.find env.vars svn in
          let save = VarMap.add (node^"/"^aid^".state") st save in
          List.fold_left (fun save (st, _) ->
              let save = 
                List.fold_left (fun save (c, _, _) -> save_expr save c) save st.until in
              aux (node, aid^"."^st.st_name^".", st.body) save)
            save states
        | AST_activate(r, _) ->
          let rec do_t save = function
            | AST_activate_if(c, t, f) ->
              let save = save_expr save c in
              let save = do_t save t in
              do_t save f
            | AST_activate_body b ->
              aux (node, b.act_id^".", b.body) save
          in
          do_t save r
      in
      List.fold_left save_eq save eqs

    in
    let consts = VarMap.filter (fun k _ -> k.[0] <> '/') env.st.save in
    { env.st with save = aux env.st.root_scope consts }
      
  (*
    init_state : prog -> id -> state
  *)
  let init_state p root =
    let (n, _) = find_node_decl p root in
    let st = {
        p = p;
        root_scope = get_root_scope p root;
        save = VarMap.empty;
        outputs = (List.map (fun (_,(n,_),_) -> n) n.ret);
    } in

    let env = { st = st; vars = Hashtbl.create 42 } in

    (* calculate constants *)
    List.iter (function
        | AST_const_decl(d, l) ->
          let cpath = "cst/" ^ d.c_name in
          Hashtbl.replace env.vars cpath (VCalc (fun () ->
            match eval_expr env ("cst", "") d.value with
            | [v] -> Hashtbl.replace env.vars cpath v
            | _ -> loc_error l error "Arity error in constant expression."))
        | _ -> ())
      p;
    List.iter (function
        | AST_const_decl(d, l) -> ignore (get_var env "cst" d.c_name)
        | _ -> ())
      p;

    reset_scope env st.root_scope;
    { st with save = Hashtbl.fold VarMap.add env.vars VarMap.empty }

  (*
    step : state -> io -> (state * io)
  *)
  let step st i =
    let vars = Hashtbl.create 10 in
    List.iter (fun (k, v) -> Hashtbl.replace vars ("/"^k) v) i;
    let env = { st = st; vars = vars } in
    
    activate_scope env st.root_scope; 
    if false then
      Hashtbl.iter (fun k v -> Format.printf "%s : %s@." k (str_repr_of_val v)) env.vars;
    let out = List.map (fun id -> id, get_var env "" id) st.outputs in
    do_weak_transitions env st.root_scope;
    extract_st env, out

  let consts p root =
    (init_state p root).save

end