open OUnit
open Types

let parse_typ s =
  let st = Stream.of_string s in
  let astpat = Parser.pat st in 
  let nodepat = Typer.typ Builtin.env astpat in
  Types.descr nodepat
;;

let cup = Tallying.CS.cup
let cap = Tallying.CS.cap
let singleton = Tallying.CS.singleton

let mk_pos (alpha,t) = singleton (Tallying.Pos (`Var alpha,parse_typ t))
let mk_neg (t,alpha) = singleton (Tallying.Neg (parse_typ t,`Var alpha))

let norm_tests = [
  "(`$A -> Bool)", "(`$B -> `$B)", 
                    cup (mk_pos ("B","Empty"))
                      (cap 
                        (mk_neg ("`$B","A"))
                        (mk_neg ("Bool","B"))
                      );
  "`$B", "`$A", mk_neg ("`$B","A");
  "`$B", "Empty", mk_pos ("B","Empty");
  "Int", "`$B", mk_neg ("Int","B");
  "Int", "(`$A | `$B)", mk_neg ("Int \\ `$B","A");
  "(Int -> Bool)", "(`$A -> `$B)",
                    cup (mk_pos ("A","Empty")) 
                      (cap 
                        (mk_pos ("A","Int"))
                        (mk_neg ("Bool","B"))
                      );
]

let m_compare = Tallying.CS.M.compare Types.compare

module MList = OUnitDiff.ListSimpleMake (struct 
  type t = Tallying.CS.m
  let compare = m_compare
  let pp_printer = Tallying.CS.print_m
  let pp_print_sep = OUnitDiff.pp_comma_separator
end)

let test_norm =
  "test tallying norm" >:::
    List.map (fun (s,t,expected) ->
      (Printf.sprintf " %s \\ %s" s t) >:: (fun _ ->
        let s = parse_typ s in
        let t = parse_typ t in
        let result = Tallying.norm (diff s t) in
        let elem s = List.sort m_compare (Tallying.CS.S.elements s) in
        MList.assert_equal (elem expected) (elem result)
      )
    ) norm_tests
;;

let mk_e l =
  List.fold_left (fun acc (v,t) -> 
    Tallying.CS.E.add (`Var v) (parse_typ t) acc
  ) Tallying.CS.E.empty l

let merge_tests = [
  [("`$A", "Empty");("`$B", "Empty")], cap (mk_pos ("A", "Empty")) (mk_pos ("B", "Empty"));
  [("`$A", "Int | Bool");("Int","`$B");("`$B", "Empty")], Tallying.CS.unsat;
  [("Bool","`$B"); ("`$B", "`$A"); ("`$A", "Empty")], Tallying.CS.unsat;
  [("Bool","`$B"); ("Int","`$B"); ("`$B","`$A"); ("`$A", "Int | Bool")],
    cap 
      (mk_neg ("`$B","A")) (
        cap 
          (mk_pos ("A", "Int | Bool")) (
            cap
            (mk_neg ("Int | Bool","B"))
            (mk_pos ("B","Int | Bool"))
          )
        );
  [("`$A", "`$B")], mk_pos ("A","`$B");
  [("`$B", "Empty")], mk_pos ("B","Empty");
]

let e_compare = Tallying.CS.E.compare Types.compare

module EList = OUnitDiff.ListSimpleMake (struct 
  type t = Tallying.CS.e
  let compare = e_compare
  let pp_printer = Tallying.CS.print_e
  let pp_print_sep = OUnitDiff.pp_comma_separator
end)

let test_merge =
  let print_test l =
    String.concat ";" (List.map (fun (s,t) -> Printf.sprintf " %s \\ %s" s t) l)
  in
  "test tallying merge" >:::
    List.map (fun (l,expected) ->
      (print_test l) >:: (fun _ ->
        let n = List.fold_left (fun acc (s,t) ->
          let s = parse_typ s in
          let t = parse_typ t in
          Tallying.CS.cap acc (Tallying.norm(diff s t))) Tallying.CS.S.empty l 
        in
        let result = Tallying.CS.S.fold (fun c acc ->
          try cup (Tallying.merge c) acc with Tallying.UnSatConstr -> acc
          ) n Tallying.CS.S.empty
        in
        let elem s = List.sort m_compare (Tallying.CS.S.elements s) in
        MList.assert_equal (elem expected) (elem result)
      )
    ) merge_tests
;;

let all =
  "all tests" >::: [
    test_norm;
    test_merge;
  ]
 
let main () =
  OUnit.run_test_tt_main all
;;
 
main ()