open Ast
open Util
open Ast_util
open Formula

open Interpret  (* used for constant evaluation ! *)


(* Transform SCADE program to logical formula. *)

(* node * prefix * equations *)
type scope = id * id * eqn ext list

type transform_data = {
    p           : Ast.prog;
    consts      : I.value VarMap.t;
    root_scope  : scope;
    (* future : the automata state *)
}

let f_and a b =
  if a = BConst true then b
    else if b = BConst true then a
      else BAnd(a, b)


(* f_of_nexpr :
    transform_data -> (string, string) -> (num_expr list -> 'a) -> expr -> 'a
*)
let rec f_of_nexpr td (node, prefix) where expr =
  let sub = f_of_nexpr td (node, prefix) in
  match fst expr with
  (* ident *)
  | AST_identifier(id, _) -> where [NIdent (node^"/"^id)]
  | AST_idconst(id, _) ->
    begin let x = VarMap.find ("cst/"^id) td.consts in
      try where [NIntConst (I.as_int x)]
      with _ -> try where [NRealConst (I.as_real x)]
      with _ -> error "Invalid data for supposedly numerical constant."
    end
  (* numerical *)
  | AST_int_const(i, _) -> where [NIntConst(int_of_string i)]
  | AST_real_const(r, _) -> where [NRealConst(float_of_string r)]
  | AST_unary(op, e) ->
    sub (function
      | [x] -> where [NUnary(op, x)]
      | _ -> invalid_arity "Unary operator") e
  | AST_binary(op, a, b) ->
    sub (function
      | [x] ->
        sub (function
          | [y] -> where [NBinary(op, x, y)]
          | _ -> invalid_arity "binary_operator") b
      | _ -> invalid_arity "binary operator") a
  (* temporal *)
  | AST_pre(_, id) -> where [NIdent id]
  | AST_arrow(a, b) ->
      BOr(
        BAnd(BRel(AST_EQ, NIdent(node^"/"^prefix^"time"), NIntConst 0),
          sub where a),
        BAnd(BRel(AST_GE, NIdent(node^"/"^prefix^"time"), NIntConst 1),
          sub where b))
  (* other *)
  | AST_if(c, a, b) ->
    BOr(
      BAnd(f_of_bexpr td (node, prefix) c, sub where a),
      BAnd(BNot(f_of_bexpr td (node, prefix) c), sub where b))
  | AST_instance ((f, _), args, nid) ->
    let (n, _) = find_node_decl td.p f in
    where (List.map (fun (_, (id, _), _) -> NIdent (node^"/"^nid^"/"^id)) n.ret)
    
  (* boolean values treated as integers *)
  | _ ->
    BOr(
      BAnd ((f_of_bexpr td (node, prefix) expr), where [NIntConst 1]),
      BAnd (BNot(f_of_bexpr td (node, prefix) expr), where [NIntConst 0])
    )

(* f_of_expr :
    transform_data -> (string, string) -> expr -> bool_expr
*)
and f_of_bexpr td (node, prefix) expr =
  let sub = f_of_bexpr td (node, prefix) in
  match fst expr with
  | AST_bool_const b -> BConst b
  | AST_binary_bool(AST_AND, a, b) -> BAnd(sub a, sub b)
  | AST_binary_bool(AST_OR, a, b) -> BOr(sub a, sub b)
  | AST_not(a) -> BNot(sub a)
  | AST_binary_rel(rel, a, b) ->
    f_of_nexpr td (node, prefix)
      (function
        | [x] -> f_of_nexpr td (node, prefix)
          (function
            | [y] -> BRel(rel, x, y)
            | _ -> invalid_arity "boolean relation") b
        | _ -> invalid_arity "boolean relation")
      a
  | _ -> loc_error (snd expr) error "Invalid type : expected boolean value"
  
  

and f_of_scope active td (node, prefix, eqs) =
  let expr_eq e eq =
    let instance_eq (id, eqs, args) =
      let eq = f_of_scope active td (node^"/"^id, "", eqs) in
      if active then
        List.fold_left (fun eq ((_,(argname,_),_), expr) ->
            let eq_arg = f_of_nexpr td (node, prefix) (function
                  | [v] -> BRel(AST_EQ, NIdent(node^"/"^id^"/"^argname), v)
                  | _ -> invalid_arity "in argument")
                expr
            in f_and eq eq_arg)
        eq args
      else
        eq
    in
    let eq = List.fold_left (fun x i -> f_and (instance_eq i) x)
      eq (extract_instances td.p e)
    in
    let pre_expr (id, expr) =
      if active then
        f_of_nexpr td (node, prefix) (function
            | [v] -> BRel(AST_EQ, NIdent("N"^id), v)
            | _ -> invalid_arity "pre on complex data not supported")
          expr
      else
        BRel(AST_EQ, NIdent("N"^id), NIdent id)
    in
    List.fold_left (fun x i -> f_and (pre_expr i) x)
      eq (extract_pre e)
  in
  let do_eq eq = match fst eq with
    | AST_assign(ids, e) ->
      let assign_eq = 
        if active then
          f_of_nexpr td (node, prefix)
            (fun vs ->
              let rels = 
                List.map2 (fun (id, _) v -> BRel(AST_EQ, NIdent (node^"/"^id), v))
                ids vs
              in
                list_fold_op f_and rels)
            e
        else
          BConst true
      in
        expr_eq e assign_eq
    | AST_assume (_, e) ->
      if active then
        f_of_bexpr td (node, prefix) e
      else
        BConst true
    | AST_guarantee _ -> BConst true
    | AST_activate (b, _) ->
      let rec cond_eq = function
        | AST_activate_body b -> BConst true
        | AST_activate_if(c, a, b) ->
          f_and (expr_eq c (BConst true))
            (f_and (cond_eq a) (cond_eq b))
      in
      let rec do_tree_act = function
        | AST_activate_body b ->
            f_of_scope true td (node, b.act_id^"_", b.body)
        | AST_activate_if(c, a, b) ->
          BOr(
            f_and (f_of_bexpr td (node, prefix) c)
              (f_and (do_tree_act a) (do_tree_inact b)),
            f_and (BNot(f_of_bexpr td (node, prefix) c))
              (f_and (do_tree_act b) (do_tree_inact a))
          )
      and do_tree_inact = function
        | AST_activate_body b ->
          f_of_scope false td (node, b.act_id^"_", b.body)
        | AST_activate_if(_, a, b) ->
          f_and (do_tree_inact a) (do_tree_inact b)
      in
        f_and (cond_eq b) (if active then do_tree_act b else do_tree_inact b)
    | AST_automaton _ -> not_implemented "f_of_scope do_eq automaton"
  in
  let time_incr_eq =
    if active then
      BRel(AST_EQ, NIdent("N"^node^"/"^prefix^"time"),
          NBinary(AST_PLUS, NIntConst 1, NIdent(node^"/"^prefix^"time")))
    else
      BRel(AST_EQ,
          NIdent("N"^node^"/"^prefix^"time"),
          NIdent(node^"/"^prefix^"time"))
  in
  List.fold_left f_and
    time_incr_eq
    (List.map do_eq eqs)

and f_of_prog p root =
    let td = {
      root_scope = get_root_scope p root;
      p = p;
      consts = I.consts p root;
    } in

    f_of_scope true td td.root_scope