open Util
open Ast
open Formula
open Ast_printer

let string_of_binary_rel = function
  | AST_EQ -> "="
  | AST_NE -> "≠"
  | AST_LT -> "<"
  | AST_LE -> "≤"
  | AST_GT -> ">"
  | AST_GE -> "≥"

let ne_prec = function
  | NUnary _ -> unary_precedence
  | NBinary(op, _, _, _) -> binary_op_precedence op
  | _ -> 100

let be_prec = function
  | BRel(op, _, _, _) -> binary_rel_precedence op
  | BAnd _ -> binary_bool_precedence AST_AND
  | BOr _ -> binary_bool_precedence AST_OR
  | BNot _ -> unary_precedence
  | _ -> 100

let is_or = function
  | BOr _ -> true
  | _ -> false
let is_and = function
  | BAnd _ -> true
  | _ -> false


let print_ch fmt pf fa a fe e =
  if fa a < fe e
    then Format.fprintf fmt "@[<hv 2>(%a)@]" pf a
    else Format.fprintf fmt "@[<hv 2>%a@]" pf a
let print_ah fmt pf fa a fe e =
  if fa a <= fe e
    then Format.fprintf fmt "@[<hv 2>(%a)@]" pf a
    else Format.fprintf fmt "@[<hv 2>%a@]" pf a

let print_id fmt id =
    let re = Str.regexp "/" in
    Format.fprintf fmt "%s" (Str.global_replace re "·" id)

let rec print_num_expr fmt e = match e with
  | NIntConst i -> Format.fprintf fmt "%d" i
  | NRealConst f -> Format.fprintf fmt "%f" f
  | NIdent id ->
    print_id fmt id
  | NBinary(op, a, b, is_real) ->
    print_ch fmt print_num_expr ne_prec a ne_prec e;
    Format.fprintf fmt "@ %s%s " (string_of_binary_op op)
        (if is_real then "." else "");
    print_ah fmt print_num_expr ne_prec b ne_prec e
  | NUnary(op, a, is_real) ->
    Format.fprintf fmt "%s%s " (string_of_unary_op op)
        (if is_real then "." else "");
    print_ah fmt print_num_expr ne_prec a ne_prec e

(* Enumeated expressions *)
let print_enum_expr fmt = function
  | EIdent id -> print_id fmt id
  | EItem x -> Format.fprintf fmt "%s" x

let str_of_enum_op = function
  | E_EQ -> "≡"
  | E_NE -> "≢"

let print_econs fmt (op, a, b) =
    Format.fprintf fmt "%a %s %a"
      print_id a (str_of_enum_op op) print_enum_expr b

(* Print boolean form of formula *)

let rec print_bool_expr fmt e = match e with
  | BConst b -> Format.fprintf fmt "%s" (if b then "true" else "false")
  | BEnumCons c -> print_econs fmt c
  | BRel(op, a, b, is_real) ->
    print_ch fmt print_num_expr ne_prec a be_prec e;
    Format.fprintf fmt "@ %s%s " (string_of_binary_rel op)
        (if is_real then "." else "");
    print_ch fmt print_num_expr ne_prec b be_prec e
  | BAnd (a, b) ->

    if is_and a then
      Format.fprintf fmt "%a" print_bool_expr a
    else
      if be_prec a < be_prec e || is_or a
        then Format.fprintf fmt "@[<hv 2>(%a)@]" print_bool_expr a
        else Format.fprintf fmt "@[<hv 2>%a@]" print_bool_expr a;
    Format.fprintf fmt "@ ∧ ";
    if is_and b then
        Format.fprintf fmt "%a" print_bool_expr b
    else
      if be_prec b < be_prec e || is_or b
        then Format.fprintf fmt "@[<hv 2>(%a)@]" print_bool_expr b
        else Format.fprintf fmt "@[<hv 2>%a@]" print_bool_expr b

  | BOr (a, b) ->
    print_ch fmt print_bool_expr be_prec a be_prec e;
    Format.fprintf fmt "@ ∨ ";
    print_ch fmt print_bool_expr be_prec b be_prec e
  | BNot (a) ->
    Format.pp_print_string fmt "¬";
    print_ch fmt print_bool_expr be_prec a be_prec e

let print_expr fmt e =
  Format.fprintf fmt "@[<hv 2>%a@]" print_bool_expr e


(* Print constraint list form of formula *)

let print_cons_sg fmt sg =
  let sg_str = match sg with
  | CONS_EQ -> "="
  | CONS_NE -> "≠"
  | CONS_GT -> ">"
  | CONS_GE -> "≥"
  in
  Format.fprintf fmt "%s" sg_str

let print_cons fmt (eq, sg) =
  Format.fprintf fmt "@[<hv 2>%a %a 0@]"
    print_num_expr eq print_cons_sg sg

let rec print_conslist fmt (econs, cons, e) =
  let c = 
    (List.map (fun x -> Left x) econs) @ (List.map (fun y -> Right y) cons)
  in
  let f fmt = function
    | Left x -> Format.fprintf fmt "@[<hv 2>%a@]" print_econs x
    | Right x -> Format.fprintf fmt "@[<hv 2>%a@]" print_cons x
  in
  let rec aux = function
    | [] -> false
    | [a] ->
      f fmt a; true
    | p::q ->
      f fmt p;
      Format.fprintf fmt "@ ∧ ";
      aux q
  in
  match e with
  | CLTrue ->
    Format.fprintf fmt "@[<hv 2>(";
    ignore (aux c);
    Format.fprintf fmt ")@]"
  | CLFalse -> 
    Format.fprintf fmt "false"
  | _ ->
    Format.fprintf fmt "@[<hv 2>(";
    if aux c then Format.fprintf fmt "@ ∧ ";
    print_conslist_expr fmt e;
    Format.fprintf fmt ")@]"

and print_conslist_expr fmt = function
  | CLTrue -> Format.fprintf fmt "true"
  | CLFalse -> Format.fprintf fmt "false"
  | CLOr(a, b) ->
    Format.fprintf fmt "@[<hv 2>(%a@ ∨ %a)@]"
       print_conslist a
        print_conslist b
  | CLAnd(a, b) ->
    Format.fprintf fmt "%a@ ∧ %a"
        print_conslist_expr a
        print_conslist_expr b