open Ast
open Data
open Util
open Ast_util

(* Data structures for the interpret *)

type scope = string

type svalue =
  | VInt of int
  | VBool of bool
  | VReal of float
  | VState of id
  | VPrevious of value
  | VBusy   (* intermediate value : calculating ! for detection of cycles *)
and value = svalue list

type state = svalue VarMap.t

(* functions for recursively getting variables *)

type calc_fun = 
  | F of (state -> calc_map -> state)
and calc_map = calc_fun VarMap.t

let get_var (st: state) (c: calc_map) (id: id) : (state * svalue) =
  let st =
    if VarMap.mem id st then st
    else match VarMap.find id c with
    | F f ->
      (* Format.printf "%s[ " id; *)
      let r = f (VarMap.add id VBusy st) c in
      (* Format.printf "]%s " id; *)
      r
  in
    match VarMap.find id st with
    | VBusy -> combinatorial_cycle id
    | v -> st, v

(* pretty-printing *)

let rec str_of_value = 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_of_value v) "" p
      ^ "]"
  | VBusy -> "#"
let print_state st =
  VarMap.iter (fun id v -> Format.printf "%s = %s@." id (str_of_value v)) st




(* Expression evaluation *)

type eval_env = {
  p: prog;
  c: calc_map;
  scopes: string list;
}

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

(* Constant calculation *)

let rec calc_const p d st c =
  let env = { p = p; c = c; scopes = ["cst"] } in
  match eval_expr env st d.value with
  | st, [v] -> VarMap.add ("cst/"^d.c_name) v st
  | _ -> type_error ("Cannot assign tuple to constant" ^ d.c_name)

let program_consts p =
  let cdecl = extract_const_decls p in
  let ccmap = List.fold_left
    (fun c (d, _) -> VarMap.add ("cst/"^d.c_name) (F (calc_const p d)) c)
    VarMap.empty cdecl
  in
  List.fold_left
    (fun st (d, _) -> let st, _ = get_var st ccmap ("cst/"^d.c_name) in st)
    VarMap.empty
    cdecl

(* get initial state for program *)
let program_init_state p root_node =
  let (n, _) = find_node_decl p root_node in
  let rec aux st p scope eqs =
    let st = VarMap.add (scope^"/init") (VBool true) st in
    let add_subscopes =
      List.fold_left (fun st (ss_id, ss_eqs, _) -> aux st p (scope^"/"^ss_id) ss_eqs)
    in
    let add_eq st eq = match fst eq with
        | AST_assign(_, e) -> add_subscopes st (extract_instances p e)
        | AST_assume _ | AST_guarantee _ -> st
        | AST_automaton (aid, astates, ret) ->
          let (initial, _) = List.find (fun (s, _) -> s.initial) astates in
          let st = VarMap.add (scope^":"^aid^".state") (VState initial.name) st in
          let add_astate st ((astate:Ast.state), _) =
            aux st p (scope^":"^aid^"."^astate.name) astate.body
          in
          List.fold_left add_astate st astates
        | AST_activate _ -> not_implemented "program_init_state activate"
    in
    List.fold_left add_eq st eqs
  in
    aux (program_consts p) p "" n.body


(* Program execution *)

(* build calc map *)
let rec build_prog_ccmap p scopes eqs st =
  let scope = List.hd scopes in
  let add_eq c eq = match fst eq with
    | AST_assign(vars, e) ->
      let calc st c =
        let env = { p = p; c = c; scopes = scopes } in
        let st, vals = eval_expr env st e in
        List.fold_left2
          (fun st (id,_) v -> VarMap.add (scope^"/"^id) v st)
          st vars vals
      in
      let c = List.fold_left (fun c (id, _) -> VarMap.add (scope^"/"^id) (F calc) c) c vars in

      let add_subscope c (ss_id, ss_eqs, ss_args) =
        let c = VarMap.merge disjoint_union c
          (build_prog_ccmap p [scope^"/"^ss_id] ss_eqs st)
        in
        let add_v c ((_, (id, _), _), eq) =
          let calc st c =
            let env = { p = p; c = c; scopes = scopes } in
            let st, vals = eval_expr env st eq in
            match vals with
            | [v] -> VarMap.add (scope^"/"^ss_id^"/"^id) v st
            | _ -> type_error "invalid arity"
          in
          VarMap.add (scope^"/"^ss_id^"/"^id) (F calc) c
        in
        List.fold_left add_v c ss_args
      in
      List.fold_left add_subscope c (extract_instances p e)
    | AST_assume _ | AST_guarantee _ -> c
    | AST_automaton _ -> not_implemented "build_prog_ccmap for automaton"
    | AST_activate _ -> not_implemented "build_prog_ccmap for activate"
  in
  List.fold_left add_eq VarMap.empty eqs

let extract_next_state active env eqs st =
  let csts = VarMap.filter
    (fun k _ -> String.length k > 4 && String.sub k 0 4 = "cst/")
    st
  in
  let rec aux active scopes eqs st nst =
    let scope = List.hd env.scopes in

    let r = VarMap.add (scope^"/init")
        (if active then VBool false
          else try VarMap.find (scope^"/init") st with Not_found -> assert false)
        nst
    in
    let add_subscopes active =
      List.fold_left (fun (st, nst) (ss_id, ss_eqs, _) ->
        aux active [scope^"/"^ss_id] ss_eqs st nst)
    in
    let add_eq (st, nst) eq = match fst eq with
      | AST_assign(vars, e) ->
        let st, nst = add_subscopes active (st, nst) (extract_instances env.p e) in
        List.fold_left (fun (st, nst) (pn, pe) ->
            let st, v = if active then eval_expr { env with scopes = scopes } st pe
              else st,
                try match VarMap.find (scope^"/"^pn) st with VPrevious x -> x | _ -> []
                with Not_found -> []
            in
            st, VarMap.add (scope^"/"^pn) (VPrevious v) nst)
          (st, nst) (extract_pre e)
      | AST_assume _ | AST_guarantee _ -> st, nst
      | AST_automaton _ -> not_implemented "extract_next_state automaton"
      | AST_activate _ -> not_implemented "extract_next_state activate"
    in
    List.fold_left add_eq (st, r) eqs
  in aux active env.scopes eqs st csts



let program_step p st inputs root_node =
  let (n, _) = find_node_decl p root_node in
  let st = List.fold_left
    (fun st (n, v) -> VarMap.add ("/"^n) v st) st inputs in

  let ccmap = build_prog_ccmap p [""] n.body st in

  let st = List.fold_left
    (fun st (_, (id, _), _) -> let st, _ = get_var st ccmap ("/"^id) in st)
    st n.ret in
  let outputs = List.map
    (fun (_, (id, _), _) -> let _, v = get_var st ccmap ("/"^id) in (id, v))
    n.ret in

  let st, next_st = extract_next_state true { p = p; scopes = [""]; c = ccmap } n.body st in

  st, outputs, next_st