(* 
    - Instanciate a program with a root node
    - Enumerate all the variables, give them a type
    - Give a type to expressions
*)

open Ast
open Ast_util
open Util

type typ =
    | TInt
    | TReal
    | TEnum of string list

let bool_true = "tt"
let bool_false = "ff"
let bool_type = TEnum [bool_true; bool_false]

let t_of_ast_t = function
    | AST_TINT -> TInt
    | AST_TREAL -> TReal
    | AST_TBOOL -> bool_type


(* probe? * variable full name * variable type *)
type var = bool * id * typ

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

type rooted_prog = {
    p              : prog;

    no_time_scope  : id -> bool; (* scopes in which not to introduce time variable *)
    init_scope     : id -> bool; (* scopes in which to introduce init variable *)

    root_node      : node_decl;
    root_scope     : scope;
    all_vars       : var list;
    const_vars     : var list;
}

(* Typing *)

(*
    type_expr_vl : prog -> var list -> string -> typ list
*)
let rec type_expr_vl p vl cvl node expr =
    let sub = type_expr_vl p vl cvl node in
    let err = loc_error (snd expr) type_error in
    match fst expr with
    (* Identifiers *)
    | AST_identifier(id, _) ->
        let _, _, t = List.find (fun (_, x, _) -> x = (node^"/"^id)) vl in [t]
    | AST_idconst(id, _) ->
        let _, _, t = List.find (fun (_, x, _) -> x = id) cvl in [t]
    (* On numerical values *)
    | AST_int_const _ -> [TInt]
    | AST_real_const _ -> [TReal]
    | AST_unary (_, se) ->
        begin match sub se with
        | [TInt] -> [TInt]
        | [TReal] -> [TReal]
        | _ -> err "Invalid argument for unary."
        end
    | AST_binary(_, a, b) ->
        begin match sub a, sub b with
        | [TInt], [TInt] -> [TInt]
        | [TReal], [TReal] -> [TReal]
        | [TInt], [TReal] | [TReal], [TInt] -> [TReal]
        | _ -> err "Invalid argument for binary."
        end
    | AST_cast(e, ty) ->
      begin match sub e, ty with
      | [x], AST_TINT -> [TInt]
      | [y], AST_TREAL -> [TReal]
      | _ -> err "Invalid arity for cast."
      end
    (* On boolean values *)
    | AST_bool_const _ -> [bool_type]
    | AST_binary_rel _ -> [bool_type] (* do not check subtypes... TODO? *)
    | AST_binary_bool _ -> [bool_type] (* the same *)
    | AST_not _ -> [bool_type] (* the same *)
    (* Temporal primitives *) 
    | AST_pre(e, _) -> sub e
    | AST_arrow(a, b) ->
      let ta, tb = sub a, sub b in
      if ta = tb then ta
        else err "Arrow does not have same type on both sides."
    (* other *)
    | AST_if(c, a, b) ->
      let ta, tb = sub a, sub b in
      if ta = tb then ta
        else err "If does not have same type on both branches."
    | AST_instance((f, _), _, _) ->
      (* do not check types of arguments... TODO? *)
      let (n, _) = find_node_decl p f in
      List.map (fun (_, _, t) -> t_of_ast_t t) n.ret
    | AST_tuple x -> List.flatten (List.map sub x)


(* type_expr : rp -> string -> expr -> typ list *)
let type_expr rp = type_expr_vl rp.p rp.all_vars rp.const_vars

let type_var tp node id =
    let _, _, t = List.find (fun (_, x, _) -> x = (node^"/"^id)) tp.all_vars in t


(* Program rooting *)


(*  decls_of_node : node_decl -> var_decl list       *)
let decls_of_node n = n.args @ n.ret @ n.var

(*  vars_in_node string -> var_decl list -> var list *)
let vars_in_node nid =
  List.map (fun (p, id, t) -> p, nid^"/"^id, t_of_ast_t t)

(* node_vars : prog -> id -> string -> var list      *)
let node_vars p f nid =
    let (n, _) = find_node_decl p f in
    vars_in_node nid (decls_of_node n)

(*
    clock_vars : rooted_prog -> scope -> var list
*)
let clock_vars rp (node, prefix, _) =
    let v =
      if not (rp.no_time_scope (node^"/"^prefix))
      then
        [false, node^"/"^prefix^"time", TInt]
      else [] in
    let v =
      if rp.init_scope (node^"/"^prefix)
        then
          (false, node^"/"^prefix^"init", bool_type)::v
        else v in
    v

(*
    extract_all_vars : rooted_prog -> scope -> var list

    Extracts all variables with names given according to
    naming convention used here and in transform.ml
*)
let rec extract_all_vars rp (node, prefix, eqs) n_vars =
    let vars_of_expr e =
      List.flatten
        (List.map
          (fun (f, id, eqs, args) ->
            let nv = node_vars rp.p f (node^"/"^id) in
            nv @ extract_all_vars rp (node^"/"^id, "", eqs) nv)
          (extract_instances rp.p e))
      @
      List.flatten
        (List.map
          (fun (id, expr) ->
              let id = node^"/"^id in
              List.mapi
                (fun i t -> false, id^"."^(string_of_int i), t)
                (type_expr_vl rp.p n_vars rp.const_vars node expr))
          (extract_pre e))
    in
    let vars_of_eq e = match fst e with
      | AST_assign(_, e) | AST_assume(_, e) -> vars_of_expr e
      | AST_guarantee((id, _), e) ->
          let gn = node^"/g_"^id in
          (false, gn, bool_type)::vars_of_expr e
      | AST_activate (b, _) ->
        let rec do_branch = function
          | AST_activate_body b ->
            let bvars = vars_in_node node b.act_locals in
            let b_scope = node, b.act_id^".", b.body in
            bvars @ clock_vars rp b_scope @
              extract_all_vars rp b_scope (bvars@n_vars)
          | AST_activate_if(c, a, b) ->
            vars_of_expr c @ do_branch a @ do_branch b
        in do_branch b
      | AST_automaton (aid, states, ret) ->
        let rst_states = List.flatten
            (List.map (fun (st, _) ->
                List.map (fun (_, (id, _), _) -> id)
                (List.filter (fun (_, _, rst) -> rst) st.until))
              states)
        in

        let do_state (st, _) =
          let tvars =
            List.flatten
              (List.map (fun (e, _, _) -> vars_of_expr e) st.until)
          in
          let st_scope = (node, aid^"."^st.st_name^".", st.body) in
          let svars = vars_in_node node st.st_locals in
          (if List.mem st.st_name rst_states then
            [false, node^"/"^aid^"."^st.st_name^".must_reset", bool_type]
          else [])
          @ svars @ tvars @ clock_vars rp st_scope @
            extract_all_vars rp st_scope (tvars@n_vars)
        in
        let st_ids = List.map (fun (st, _) -> st.st_name) states in
        (false, node^"/"^aid^".state", TEnum st_ids)::
        (false, node^"/"^aid^".next_state", TEnum st_ids)::
        (List.flatten (List.map do_state states))
    in
    let v = List.flatten (List.map vars_of_eq eqs) in
    v


(*
    root_prog : prog -> id -> string list -> string list -> rooted_prog
*)
let root_prog p root no_time_scope init_scope =
    let (root_node, _) =
      try find_node_decl p root
      with Not_found -> error ("No such root node: " ^ root) in
    let root_scope = ("", "", root_node.body) in

    let const_vars = List.map
        (fun (cd, _) -> (false, cd.c_name, t_of_ast_t cd.typ))
        (extract_const_decls p)
    in


    let rp = {
      p; root_scope; root_node;

      no_time_scope;
      init_scope;

      const_vars;
      all_vars = [] } in

    let root_vars = vars_in_node "" (decls_of_node root_node) in
    let root_clock_vars = clock_vars rp root_scope in

    { rp with all_vars = root_clock_vars @
            root_vars @ extract_all_vars rp root_scope root_vars }