open Ast
open Util
open Ast_util
open Formula
open Typing

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 = {
    rp          : rooted_prog;
    consts      : I.value VarMap.t;
    (* future : the automata state *)
}

(* 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(expr, id) ->
    let typ = type_expr td.rp node expr in
      where (List.mapi (fun i _ -> NIdent (id^"."^(string_of_int i))) typ)
  | AST_arrow(a, b) ->
    f_or
      (f_and (BRel(AST_EQ, NIdent(node^"/"^prefix^"time"), NIntConst 0))
        (sub where a))
      (f_and (BRel(AST_GE, NIdent(node^"/"^prefix^"time"), NIntConst 1))
        (sub where b))
  (* other *)
  | AST_if(c, a, b) ->
    f_or
      (f_and (f_of_expr td (node, prefix) c) (sub where a))
      (f_and (BNot(f_of_expr td (node, prefix) c)) (sub where b))
  | AST_instance ((f, _), args, nid) ->
    let (n, _) = find_node_decl td.rp.p f in
    where (List.map (fun (_, id, _) -> NIdent (node^"/"^nid^"/"^id)) n.ret)
    
  (* boolean values treated as integers *)
  | _ ->
    f_or
      (f_and ((f_of_expr td (node, prefix) expr)) (where [NIntConst 1]))
      (f_and (BNot(f_of_expr td (node, prefix) expr)) (where [NIntConst 0]))


(* 
  f_of_expr :
    transform_data -> (string, string) -> expr -> bool_expr
  f_of_bexpr :
    transform_data -> (string, string) -> (bool_expr -> 'a) -> expr -> 'a
*)
and f_of_bexpr td (node, prefix) where expr =
  let sub = f_of_bexpr td (node, prefix) in
  match fst expr with
  | AST_bool_const b -> where (BConst b)
  | AST_binary_bool(AST_AND, a, b) -> f_and (sub where a) (sub where b)
  | AST_binary_bool(AST_OR, a, b) -> f_or (sub where a) (sub where b)
  | AST_not(a) -> BNot(sub where a)
  | AST_binary_rel(rel, a, b) ->
    where
      (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)
  | AST_identifier (id, _) ->
    let id = node^"/"^id in
    f_or
      (f_and (BBoolEq(id, true)) (where (BConst true)))
      (f_and (BBoolEq(id, false)) (where (BConst false)))
  (* temporal *)
  | AST_pre(_, id) ->
    f_or
      (f_and (BBoolEq(id, true)) (where (BConst true)))
      (f_and (BBoolEq(id, false)) (where (BConst false)))
  | AST_arrow(a, b) ->
      f_or
        (f_and (BRel(AST_EQ, NIdent(node^"/"^prefix^"time"), NIntConst 0))
          (sub where a))
        (f_and (BRel(AST_GE, NIdent(node^"/"^prefix^"time"), NIntConst 1))
          (sub where b))
  (* other *)
  | AST_if(c, a, b) ->
    f_or
      (f_and (f_of_expr td (node, prefix) c) (sub where a))
      (f_and (BNot(f_of_expr td (node, prefix) c)) (sub where b))

  | _ -> loc_error (snd expr) error "Expected boolean value."

and f_of_expr td (node, prefix) expr =
    f_of_bexpr td (node, prefix) (fun x -> x) expr
  

(*
    Translate program into one big formula
*)
let rec f_of_scope active td (node, prefix, eqs) assume_guarantees =
  let expr_eq e eq =
    let instance_eq (_, id, eqs, args) =
      let eq = f_of_scope active td (node^"/"^id, "", eqs) assume_guarantees 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.rp.p e)
    in
    let pre_expr (id, expr) =
      if active then
        f_of_nexpr td (node, prefix) (fun elist ->
            list_fold_op f_and
              (List.mapi
                (fun i v -> BRel(AST_EQ, NIdent("N"^id^"."^(string_of_int i)), v))
                elist))
          expr
      else
        let typ = type_expr td.rp node expr in
        list_fold_op f_and
          (List.mapi
            (fun i _ -> let x = string_of_int i in BRel(AST_EQ, NIdent("N"^id^"."^x), NIdent (id^"."^x)))
            typ)
    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) ->
      let assume_eq = 
        if active then
          f_of_expr td (node, prefix) e
        else
          BConst true
      in
        expr_eq e assume_eq
    | AST_guarantee (_, e) ->
      let guarantee_eq =
        if active && assume_guarantees then
          f_of_expr td (node, prefix) e
        else
          BConst true
      in
        expr_eq e guarantee_eq
    | 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) assume_guarantees
        | AST_activate_if(c, a, b) ->
          f_or
            (f_and (f_of_expr td (node, prefix) c)
              (f_and (do_tree_act a) (do_tree_inact b)))
            (f_and (BNot(f_of_expr 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) assume_guarantees
        | 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 rp assume_guarantees =
    let td = {
      rp = rp;
      consts = I.consts rp;
    } in

    f_of_scope true td td.rp.root_scope assume_guarantees

(*
    Translate init state into a formula
*)
let rec init_f_of_scope rp (node, prefix, eqs) =
  let expr_eq e =
    let instance_eq (_, id, eqs, args) =
     init_f_of_scope rp (node^"/"^id, "", eqs)
    in
    List.fold_left (fun x i -> f_and (instance_eq i) x)
      (BConst true) (extract_instances rp.p e)
  in
  let do_eq eq = match fst eq with
    | AST_assign(_, e) | AST_assume(_, e) | AST_guarantee(_, e) ->
        expr_eq e
    | AST_activate (b, _) ->
      let rec cond_eq = function
        | AST_activate_body b ->
          init_f_of_scope rp (node, b.act_id^".", b.body)
        | AST_activate_if(c, a, b) ->
          f_and (expr_eq c)
            (f_and (cond_eq a) (cond_eq b))
      in
      cond_eq b
    | AST_automaton _ -> not_implemented "f_of_scope do_eq automaton"
  in
  let time_eq =
    BRel(AST_EQ,
        NIdent(node^"/"^prefix^"time"),
        NIntConst 0)
  in
  List.fold_left f_and time_eq (List.map do_eq eqs)

and init_f_of_prog rp =
    init_f_of_scope rp rp.root_scope

(*
    Get expressions for guarantee violation
*)
let rec g_of_scope td (node, prefix, eqs) cond =
  let expr_g e =
    let instance_g (_, id, eqs, args) =
        g_of_scope td (node^"/"^id, "", eqs) cond
    in
    List.fold_left (fun x i -> (instance_g i) @ x)
      [] (extract_instances td.rp.p e)
  in
  let do_eq eq = match fst eq with
    | AST_assign(_, e) | AST_assume(_, e) ->
        expr_g e
    | AST_guarantee((id, _), e) ->
        (id, f_and cond (f_of_expr td (node, prefix) (AST_not(e), snd e)))
          :: (expr_g e)
    | AST_activate (b, _) ->
      let rec cond_g cond = function
        | AST_activate_body b ->
          g_of_scope td (node, b.act_id^".", b.body) cond
        | AST_activate_if(c, a, b) ->
          (cond_g (f_and cond (f_of_expr td (node, prefix) c)) a) @
          (cond_g (f_and cond (f_of_expr td (node, prefix) (AST_not(c), snd c))) b) @
          (expr_g c)
      in
      cond_g cond b
    | AST_automaton _ -> not_implemented "g_of_scope automaton"
  in
  List.fold_left (@) [] (List.map do_eq eqs)

and guarantees_of_prog rp =
    let td = {
      rp = rp;
      consts = I.consts rp;
    } in

    g_of_scope td rp.root_scope (BConst true)