open Ast
open Util
open Ast_util

(* AST for logical formulas *)


(* Numerical part *)

(* bool on numerical operators : is it float ? *)

type num_expr =
  (* constants *)
  | NIntConst of int
  | NRealConst of float
  (* operators *)
  | NBinary of binary_op * num_expr * num_expr * bool
  | NUnary of unary_op * num_expr * bool
  (* identifier *)
  | NIdent of id

type num_cons_op =
  | CONS_EQ | CONS_NE
  | CONS_GT | CONS_GE
type num_cons = num_expr * num_cons_op  (* always imply right member = 0 *)


(* Logical part *)


(* Enumerated types *)
type enum_expr =
  | EIdent of id
  | EItem of string

type enum_op =
  | E_EQ | E_NE

type enum_cons = enum_op * id * enum_expr

type bool_expr =
  (* constants *)
  | BConst of bool
  (* operators from numeric values to boolean values *)
  | BRel of binary_rel_op * num_expr * num_expr * bool
  (* operators on enumerated types *)
  | BEnumCons of enum_cons
  (* boolean operators *)
  | BAnd of bool_expr * bool_expr
  | BOr of bool_expr * bool_expr
  | BTernary of bool_expr * bool_expr * bool_expr
      (* (A and B) or ((not A) and C) *)
  | BNot of bool_expr

let is_false = function
  | BConst false | BNot(BConst true) -> true
  | _ -> false
let is_true = function
  | BConst true | BNot(BConst false) -> true
  | _ -> false

let f_and a b =
  if is_false a || is_false b then BConst false
  else if is_true a then b
  else if is_true b then a
  else BAnd(a, b)

let f_and_list = List.fold_left f_and (BConst true)

let f_or a b =
  if is_true a || is_true b then BConst true
  else if is_false a then b
  else if is_false b then a
  else BOr(a, b)

let f_ternary c a b =
  if is_true c then a
  else if is_false c then b
  else if is_true a && is_false b then c
  else if is_true b && is_false a then BNot c
  else BTernary(c, a, b)

let f_e_op op a b = match a, b with
  | EItem i, EItem j -> BConst (if op = E_EQ then i = j else i <> j)
  | EIdent x, v | v, EIdent x -> BEnumCons(op, x, v)
let f_e_eq = f_e_op E_EQ

(* Write all formula without using the NOT operator *)

let rec eliminate_not = function
  | BNot e -> eliminate_not_negate e
  | BAnd(a, b) -> BAnd(eliminate_not a, eliminate_not b)
  | BOr(a, b) -> BOr(eliminate_not a, eliminate_not b)
  | BTernary(c, a, b) ->
    BTernary(eliminate_not c, eliminate_not a, eliminate_not b)
  | x -> x
and eliminate_not_negate = function
  | BConst x -> BConst(not x)
  | BEnumCons(op, a, b) -> BEnumCons((if op = E_EQ then E_NE else E_EQ), a, b)
  | BNot e -> eliminate_not e
  | BRel(r, a, b, is_real) ->
    if r = AST_EQ then
      BOr(BRel(AST_LT, a, b, is_real), BRel(AST_GT, a, b, is_real))
    else
      let r' = match r with
      | AST_EQ -> AST_NE
      | AST_NE -> AST_EQ
      | AST_LT -> AST_GE
      | AST_LE -> AST_GT
      | AST_GT -> AST_LE
      | AST_GE -> AST_LT
      in
      BRel(r', a, b, is_real)
  | BTernary (c, a, b) -> 
    eliminate_not_negate(BOr(BAnd(c, a), BAnd(BNot c, b)))
  | BAnd(a, b) ->
    BOr(eliminate_not_negate a, eliminate_not_negate b)
  | BOr(a, b) ->
    BAnd(eliminate_not_negate a, eliminate_not_negate b)



(*
  In big ANDs, try to separate levels of /\ and levels of \/
  We also use this step to simplify trues and falses that may be present.
*)

type conslist = enum_cons list * num_cons list * conslist_bool_expr
and conslist_bool_expr =
  | CLTrue
  | CLFalse
  | CLAnd of conslist_bool_expr * conslist_bool_expr
  | CLOr of conslist * conslist

let rec conslist_of_f = function
  | BNot e -> conslist_of_f (eliminate_not_negate e)
  | BRel (op, a, b, is_real) ->
    let x, y, op = match op with
      | AST_EQ -> a, b, CONS_EQ
      | AST_NE -> a, b, CONS_NE
      | AST_GT -> a, b, CONS_GT
      | AST_GE -> a, b, CONS_GE
      | AST_LT -> b, a, CONS_GT
      | AST_LE -> b, a, CONS_GE
    in
    let cons = if y = NIntConst 0 || y = NRealConst 0.
      then (x, op)
      else (NBinary(AST_MINUS, x, y, is_real), op)
    in [], [cons], CLTrue
  | BConst x ->
    [], [], if x then CLTrue else CLFalse
  | BEnumCons e ->
    [e], [], CLTrue
  | BOr(a, b) ->
    let eca, ca, ra = conslist_of_f a in
    let ecb, cb, rb = conslist_of_f b in
    begin match eca, ca, ra, ecb, cb, rb with
      | _, _, CLFalse, _, _, _ -> ecb, cb, rb
      | _, _, _, _, _, CLFalse -> eca, ca, ra
      | [], [], CLTrue, _, _, _ -> [], [], CLTrue
      | _, _, _, [], [], CLTrue -> [], [], CLTrue
      | _ -> [], [], CLOr((eca, ca, ra), (ecb, cb, rb))
    end
  | BAnd(a, b) ->
    let eca, ca, ra = conslist_of_f a in
    let ecb, cb, rb = conslist_of_f b in
    let cons = ca @ cb in
    let econs = eca @ ecb in
    begin match ra, rb with
      | CLFalse, _ | _, CLFalse -> [], [], CLFalse
      | CLTrue, _ -> econs, cons, rb
      | _, CLTrue -> econs, cons, ra
      | _, _ -> econs, cons, CLAnd(ra, rb)
    end
  | BTernary(c, a, b) ->
    conslist_of_f (BOr (BAnd(c, a), BAnd(BNot(c), b)))



(*
  Simplify formula considering something is true
*)

let rec get_root_true = function
  | BAnd(a, b) ->
    get_root_true a @ get_root_true b
  | BEnumCons e -> [BEnumCons e]
  | BRel (a, b, c, d) -> [BRel (a, b, c, d)]
  | _ -> []

let rec simplify true_eqs e =
  if List.exists
      (fun f ->
        try eliminate_not e = eliminate_not f with _ -> false)
      true_eqs
    then BConst true
  else if List.exists
      (fun f ->
        match e, f with
        | BEnumCons(E_EQ, a, EItem v), BEnumCons(E_EQ, b, EItem w)
          when a = b && v <> w -> true
        | _ -> try eliminate_not e = eliminate_not_negate f with _ -> false)
      true_eqs
    then BConst false
  else
    match e with
    | BAnd(a, b) ->
      f_and
        (simplify true_eqs a)
        (simplify true_eqs b)
    | BOr(a, b) ->
      f_or
        (simplify true_eqs a)
        (simplify true_eqs b)
    | BTernary(c, a, b) ->
      f_ternary
        (simplify true_eqs c)
        (simplify true_eqs a)
        (simplify true_eqs b)
    | BNot(n) ->
      begin match simplify true_eqs n with
      | BConst e -> BConst (not e)
      | x -> BNot x
      end
    | v -> v
    
let rec simplify_k true_eqs e =
  List.fold_left f_and (simplify true_eqs e) true_eqs
  


(*
  Simplify a formula replacing a variable with another
*)

let rec formula_replace_evars repls e =
  let new_name x =
    try List.assoc x repls
    with Not_found -> x
  in
  match e with
  | BOr(a, b) ->
    f_or (formula_replace_evars repls a) (formula_replace_evars repls b)
  | BAnd(a, b) ->
    f_and (formula_replace_evars repls a) (formula_replace_evars repls b)
  | BTernary(c, a, b) ->
    f_ternary
      (formula_replace_evars repls c)
      (formula_replace_evars repls a)
      (formula_replace_evars repls b)
  | BNot(n) ->
    begin match formula_replace_evars repls n with
    | BConst e -> BConst (not e)
    | x -> BNot x
    end
  | BEnumCons (op, a, EIdent b) ->
    let a', b' = new_name a, new_name b in
    if a' = b' then
      match op with | E_EQ -> BConst true | E_NE -> BConst false
    else
      BEnumCons(op, a', EIdent b')
  | BEnumCons (op, a, EItem x) ->
    BEnumCons (op, new_name a, EItem x)
  | x -> x



(*
  extract names of variables referenced in formula
*)

let rec refd_evars_of_f = function
  | BAnd (a, b) | BOr(a, b) -> refd_evars_of_f a @ refd_evars_of_f b
  | BNot a -> refd_evars_of_f a
  | BTernary(a, b, c) -> refd_evars_of_f a @ refd_evars_of_f b @ refd_evars_of_f c
  | BEnumCons(_, e, EItem _) -> [e]
  | BEnumCons(_, e, EIdent f) -> [e; f]
  | _ -> []