(* 
    - 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 = "T"
let bool_false = "F"
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;
    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
    (* 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


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


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

(*
    extract_all_vars : prog -> scope -> var 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) n_vars const_vars =
    let vars_of_expr e =
      List.flatten
        (List.map
          (fun (f, id, eqs, args) ->
            let nv = node_vars p f (node^"/"^id) in
            nv @ extract_all_vars p (node^"/"^id, "", eqs) nv const_vars)
          (extract_instances p e))
      @
      List.flatten
        (List.map
          (fun (id, expr) ->
              let vd =
                List.mapi
                  (fun i t -> false, id^"."^(string_of_int i), t)
                  (type_expr_vl p n_vars const_vars node expr)
              in
              vd @ (List.map (fun (a, x, c) -> (a, "N"^x, c)) vd))
          (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 ->
            let bvars = vars_in_node node b.act_locals in
            bvars @
              extract_all_vars p
                (node, b.act_id^".", b.body)
                (bvars@n_vars)
                const_vars
          | 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", TInt)::
      (false, "N"^node^"/"^prefix^"time", TInt)::
      (List.flatten (List.map vars_of_eq eqs))


(*
    root_prog : prog -> id -> rooted_prog
*)
let root_prog p root =
    let (n, _) = find_node_decl p root in
    let root_scope = ("", "", n.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 root_vars = vars_in_node "" (decls_of_node n) in
    
    { p; root_scope; root_node = n;
      const_vars = const_vars;
      all_vars = root_vars @ extract_all_vars p root_scope root_vars const_vars }