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

(* Numerical types / Enumerated types *)
type ne_expr =
  | EE of enum_expr
  | NE of num_expr

(* f_of_neexpr :
    transform_data -> (string, string) -> (ne_expr list -> bool_expr) -> expr -> bool_expr
*)
let rec f_of_neexpr td (node, prefix) where expr =
  let sub = f_of_neexpr td (node, prefix) in
  let le = loc_error (snd expr) in
  match fst expr with
  (* ident *)
  | AST_identifier(id, _) ->
    let qid = node^"/"^id in
    begin match type_var td.rp node id with
    | TInt | TReal -> where [NE (NIdent qid)]
    | TEnum _ -> where [EE (EIdent qid)]
    end
  | AST_idconst(id, _) ->
    begin let x = VarMap.find ("cst/"^id) td.consts in
      try where [NE (NIntConst (I.as_int x))]
      with _ -> try where [NE (NRealConst (I.as_real x))]
      with _ -> try where [EE (EItem (if I.as_bool x then bool_true else bool_false))]
      with _ -> le error "Invalid data for supposedly numerical/boolean constant."
    end
  (* numerical *)
  | AST_int_const(i, _) -> where [NE(NIntConst(int_of_string i))]
  | AST_real_const(r, _) -> where [NE(NRealConst(float_of_string r))]
  | AST_bool_const b -> where [EE(EItem (if b then bool_true else bool_false))]
  | AST_unary(op, e) ->
    sub (function
      | [NE x] -> where [NE(NUnary(op, x))]
      | _ -> le invalid_arity "Unary operator") e
  | AST_binary(op, a, b) ->
    sub (function
      | [NE x] ->
        sub (function
          | [NE y] -> where [NE(NBinary(op, x, y))]
          | _ -> le invalid_arity "binary operator") b
      | _ -> le invalid_arity "binary operator") a
  (* temporal *)
  | AST_pre(expr, id) ->
    let typ = type_expr td.rp node expr in
      where
        (List.mapi
          (fun i t -> let x = id^"."^(string_of_int i) in
              match t with
              | TInt | TReal -> NE(NIdent x)
              | TEnum _ -> EE(EIdent x))
          typ)
  | AST_arrow(a, b) ->
    if td.rp.init_scope (node^"/")
    then
      f_or
        (f_and (f_e_eq (EIdent (node^"/"^prefix^"init")) (EItem bool_true))
          (sub where a))
        (f_and (f_e_eq (EIdent (node^"/"^prefix^"init")) (EItem bool_false))
          (sub where b))
    else
      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, t) -> let x = node^"/"^nid^"/"^id in
            match t with
            | AST_TINT | AST_TREAL -> NE(NIdent x)
            | _ -> EE(EIdent x))
        n.ret)
  | AST_tuple l ->
    let rec rl l x = match l with
      | [] -> where x
      | p::q -> 
        sub (fun y -> rl q (x@y)) p
    in rl l []
  (* boolean values treated as enumerated types *)
  | _ when type_expr td.rp node expr = [bool_type] ->
    f_or
      (f_and (f_of_expr td (node, prefix) expr) (where [EE (EItem bool_true)]))
      (f_and (f_of_expr td (node, prefix) (AST_not(expr), snd expr))
        (where [EE (EItem bool_false)]))
  | _ -> le type_error "Expected numerical/enumerated value"



(* 
  f_of_expr :
    transform_data -> (string, string) -> expr -> bool_expr
  f_of_bexpr :
    transform_data -> (string, string) -> (bool_expr -> bool_expr) -> expr -> bool_expr
*)
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_neexpr td (node, prefix)
        (function
          | [NE x; NE y] -> BRel(rel, x, y)
          | [EE x; EE y] ->
            let eop = match rel with
              | AST_EQ -> E_EQ
              | AST_NE -> E_NE
              | _ -> type_error "Invalid operator on enumerated values."
            in f_e_op eop x y
          | _ -> invalid_arity "Binary operator")
        (AST_tuple [a; b], snd expr))
  (* Temporal *)
  | 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))
  (* Enumerations... *)
  | _ when type_expr td.rp node expr = [bool_type] ->
    let ff = function
      | [EE x] -> 
        f_or
          (f_and (f_e_eq x (EItem bool_true)) (where (BConst true)))
          (f_and (f_e_eq x (EItem bool_false)) (where (BConst false)))
      | _ -> assert false
    in
    f_of_neexpr td (node, prefix) ff expr
  | _ -> type_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 =
    let instance_eq (_, id, eqs, args) =
      let eq = f_of_scope active td (node^"/"^id, "", eqs) assume_guarantees in
      if active then
        let arg_eq ((_,argname,_), expr) =
            f_of_neexpr td (node, prefix) (function
                | [NE v] -> BRel(AST_EQ, NIdent(node^"/"^id^"/"^argname), v)
                | [EE v] -> f_e_eq (EIdent(node^"/"^id^"/"^argname)) v
                | _ -> invalid_arity "in argument")
              expr
        in f_and eq (f_and_list (List.map arg_eq args))
      else
        eq
    in
    let eq_i = f_and_list (List.map instance_eq (extract_instances td.rp.p e)) in

    let pre_expr (id, expr) =
      if active then
        f_of_neexpr td (node, prefix) (fun elist ->
            list_fold_op f_and
              (List.mapi
                (fun i v -> let x = ("N"^id^"."^(string_of_int i)) in
                  match v with
                  | NE v -> BRel(AST_EQ, NIdent x, v)
                  | EE v -> f_e_eq (EIdent x) v)
                elist))
          expr
      else
        let typ = type_expr td.rp node expr in
        list_fold_op f_and
          (List.mapi
            (fun i t -> let x = string_of_int i in
              match t with
              | TInt | TReal -> BRel(AST_EQ, NIdent("N"^id^"."^x), NIdent (id^"."^x))
              | TEnum _ -> f_e_eq (EIdent("N"^id^"."^x)) (EIdent (id^"."^x)))
            typ)
    in
    let eq_p = f_and_list (List.map pre_expr (extract_pre e)) in

    f_and eq_i eq_p
  in
  let do_eq eq = match fst eq with
    | AST_assign(ids, e) ->
      let assign_eq = 
        if active then
            let apply_f vs =
              let rels = 
                List.map2 (fun (id, _) -> function
                  | NE v -> BRel(AST_EQ, NIdent (node^"/"^id), v)
                  | EE v -> f_e_eq (EIdent (node^"/"^id)) v)
                ids vs
              in
                f_and_list rels
            in
            f_of_neexpr td (node, prefix) apply_f e
        else
          BConst true
      in
        f_and (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
        f_and (expr_eq e) assume_eq
    | AST_guarantee ((id, _), e) ->
      let gn = node^"/g_"^id in
      let guarantee_eq =
        if active && assume_guarantees then
          f_and (f_of_expr td (node, prefix) e)
                (BEnumCons(E_EQ, gn, EItem bool_true))
        else
          f_or
            (f_and (f_of_expr td (node, prefix) e)
                (BEnumCons(E_EQ, gn, EItem bool_true)))
            (f_and (BNot (f_of_expr td (node, prefix) e))
                (BEnumCons(E_EQ, gn, EItem bool_false)))
      in
        f_and (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)
            (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 (aid, states, _) ->
      let stv = node^"/"^aid^".state" in
      let nstv = "N"^node^"/"^aid^".state" in
      let st_eq_inact (st, _) =
        f_and
          (f_of_scope false td
              (node, aid^"."^st.st_name^".", st.body) assume_guarantees)
          (f_and_list (List.map (fun (c, _, _) -> expr_eq c) st.until))
      in
      if active then
        let st_eq_act (st, l) =
          let act_eq =
            let st_eq = f_of_scope true td
              (node, aid^"."^st.st_name^".", st.body) assume_guarantees in
            let rec aux = function
              | [] -> BEnumCons(E_EQ, nstv, EItem st.st_name)
              | (c, (target, l), rst)::q ->
                if rst then loc_error l error "Resetting transitions not supported.";
                f_or
                  (f_and (f_of_expr td (node, prefix) c)
                        (BEnumCons(E_EQ, nstv, EItem target)))
                  (f_and (BNot (f_of_expr td (node, prefix) c)) (aux q))
            in f_and st_eq (aux st.until) 
          in
          f_or
            (f_and (BEnumCons(E_EQ, stv, EItem st.st_name)) act_eq)
            (f_and (BEnumCons(E_NE, stv, EItem st.st_name))
              (st_eq_inact (st, l)))
        in
        f_and_list (List.map st_eq_act states)
      else
        f_and_list (List.map st_eq_inact states)
  in
  let time_incr_eq =
    if active then
      f_and
        (if not (td.rp.no_time_scope (node^"/"))
          then BRel(AST_EQ, NIdent("N"^node^"/"^prefix^"time"),
              NBinary(AST_PLUS, NIntConst 1, NIdent(node^"/"^prefix^"time")))
          else BConst true)
        (if td.rp.init_scope (node^"/")
          then f_e_eq (EIdent("N"^node^"/"^prefix^"init")) (EItem bool_false)
          else BConst true)
    else
      f_and
        (if not (td.rp.no_time_scope (node^"/"))
          then BRel(AST_EQ,
              NIdent("N"^node^"/"^prefix^"time"),
              NIdent(node^"/"^prefix^"time"))
          else BConst true)
        (if td.rp.init_scope (node^"/")
          then f_e_eq (EIdent("N"^node^"/"^prefix^"init"))
                (EIdent (node^"/"^prefix^"init"))
          else BConst true)
  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 (aid, states, _) ->
      let (st, _) = List.find (fun (st, _) -> st.initial) states in
      let init_eq = BEnumCons(E_EQ, node^"/"^aid^".state", EItem st.st_name) in
      let state_eq (st, _) =
        let sc_f =
          init_f_of_scope rp (node, aid^"."^st.st_name^".", st.body)
        in List.fold_left (fun f (c, _, _) -> f_and f (expr_eq c)) sc_f st.until
      in
      List.fold_left f_and init_eq
        (List.map state_eq states)
  in
  let time_eq =
    f_and
      (if not (rp.no_time_scope (node^"/"))
        then
          BRel(AST_EQ,
            NIdent(node^"/"^prefix^"time"),
            NIntConst 0)
        else BConst true)
      (if rp.init_scope (node^"/")
        then
          f_e_eq (EIdent(node^"/"^prefix^"init")) (EItem bool_true)
        else BConst true)
  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 (aid, states, _) ->
      let st_g (st, _) =
        g_of_scope td (node, aid^"."^st.st_name^".", st.body)
          (f_and cond (BEnumCons(E_EQ, node^"/"^aid^".state", EItem st.st_name)))
      in
      List.flatten (List.map st_g states)
  in
  List.flatten (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)