open Ast
open Util
open Ast_util
open Interface

module I : INTERPRET = 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 -> value)

  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 *)

  (* path to node * subscope prefix in node * equations in scope *)
  type scope = id * id * eqn ext list

  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 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 ()
      | 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.
  *)
  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; 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 i () =
          let vals = eval_expr env (node, prefix) e in
          List.iter2 (fun (id, _) v ->
            Hashtbl.replace env.vars (node^"/"^id) v) vars vals;
          List.nth vals i
        in
        List.iteri
          (fun i (id, _) ->
              Hashtbl.replace env.vars (node^"/"^id) (VCalc (calc i)))
          vars
      | AST_assume(_,e) | AST_guarantee(_,e) -> do_expr e
      | _ -> not_implemented "activate_scope"
    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 do_weak_transitions env (node, prefix, eqs) =
    not_implemented "do_weak_transitions"

  (*
    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 =
        let save = List.fold_left
          (fun save (id, expr) ->
            VarMap.add
              (node^"/"^prefix^id)
              (VPrevious (eval_expr env (node, prefix) expr))
              save)
          save (extract_pre e) in
        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, _) -> not_implemented "save_eq automaton"
        | AST_activate(r, _) -> not_implemented "save_eq activate"
      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 root_scope = ("", "", n.body) in

    let st = {
        p = p;
        root_scope = root_scope;
        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; 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 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; 
    (*
      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


end