open Hw2_unify;;
open Printf

(* Write in console to run:                                          *)
(* ocamlc -o <EXECUTOR-NAME> -I ../HW1  ../HW1/hw1.mli ../HW1/hw1.ml  hw2_unify.mli hw2_unify.ml t_unify.ml *)
(* > ./<EXECUTOR-NAME>                                               *)
type substitution = (string * algebraic_term) list
type equation = algebraic_term * algebraic_term
type system_of_equations = equation list

let eqt11 = Fun("x", [Var "p1"]), Fun("x", [Var "p2"])
let eqt12 = Fun("y", [Var "p2"]), Fun("y", [Var "p4"])
let eqt13 = Fun("z", [Var "p5"]), Fun("z", [Var "p6"])

let eqt21 = Fun("x", [Var "p1"]), Fun("x", [Var "p2"])
let eqt22 = Fun("m", [Var "p1"]), Fun("y", [Var "p4"])
let eqt23 = Fun("z", [Var "p5"]), Fun("z", [Var "p6"])

let eqt31 = Fun("x", [Var "p1"]), Fun("x", [Var "p2"])
let eqt32 = Fun("y", [Var "p1"]), Fun("y", [Var "p4"])
let eqt33 = Fun("z", [Var "p1"]), Fun("z", [Var "p6"])

let sym1 = [eqt11; eqt12; eqt13];;
let sym2 = [eqt21; eqt22; eqt23];;
let sym3 = [eqt31; eqt32; eqt33];;


let vx = Hw1.Var "x"
let vy = Hw1.Var "y"
let vz = Hw1.Var "z"
let s = Hw1.Abs ("x", Hw1.Abs("y", Hw1.Abs ("z", Hw1.App (Hw1.App (vx, vz), Hw1.App (vy, vz)))))
let k = Hw1.Abs ("a", Hw1.Abs ("b", Hw1.Var "a"))

let lambda_to_system term =
    let cnt = ref 0 in
    let arrow l r = Fun ("A", [l; r]) in
    let rec helper t =
        match t with
        | Hw1.Var x ->
            ([], Hw2_unify.Var ("t" ^ x))
        | Hw1.App (p, q) -> cnt := !cnt + 1;
                        let new_type = Hw2_unify.Var ("t" ^ (string_of_int !cnt)) in
                        let (sys_p, type_p) = helper p in
                        let (sys_q, type_q) = helper q in
                        ((type_p, arrow type_q new_type)
                            :: (List.rev_append  sys_p sys_q), new_type)
        | Hw1.Abs (x, body) ->
                        let res = helper body in
                        (fst res, arrow (Hw2_unify.Var ("t" ^ x)) (snd res))
    in let (sys, tp) = helper term in
    (sys, tp) ;;


let subt = [("p1", Var("p3"));("p3", Var("p4"));("p5", Var("p6"))]

let rec algebraic_term_to_string (at : algebraic_term) =
    let rec impl a =
        match a with
        | Var x -> x
        | Fun(f, l) -> f ^ "(" ^ impl_for_list l ^ ")"

    and impl_for_list lt =
        match lt with
        | [] -> ""
        | (h::[]) -> impl h
        | (h::t) -> (impl h) ^ " " ^ (impl_for_list t)
    in
    impl at;;

let eqt_to_string (eqt : equation) =
    match eqt with
    | (l, r) -> algebraic_term_to_string l ^ " = " ^ algebraic_term_to_string r;;

let rec print_substitution sub =
    match sub with
    | [] -> print_string "\n"
    | (h::t) -> print_string ((fst h) ^ " = " ^ (algebraic_term_to_string (snd h)) ^ "\n"); print_substitution t;;

let print_ans ans =
    match ans with
    | Some s -> print_substitution s
    | None -> print_string "Not substitution!\n";;

let (sys, _) = lambda_to_system (Hw1.App (s, k)) in
let ans = solve_system sys in
let correct = match solve_system sys with
    | Some sub -> check_solution sub sys
    | None -> printf "No solution\n"; true in
printf "System\n";
List.iter (fun eq -> printf "%s\n" (eqt_to_string eq)) sys;
printf "Answer\n";
print_ans ans;
printf "S K : %B\n" correct;;


let maj_eqt = system_to_equation(sym1);;

let eql = apply_substitution subt (fst maj_eqt);;
let eqr = apply_substitution subt (snd maj_eqt);;

print_string ((algebraic_term_to_string eql) ^ "\n");;
print_string ((algebraic_term_to_string eqr) ^ "\n");;

let check_subt = check_solution subt sym1;;
print_string ((string_of_bool check_subt) ^ "\n");;

let ans_sym1 = solve_system sym1;;
let ans_sym2 = solve_system sym2;;
let ans_sym3 = solve_system sym3;;

print_ans ans_sym1;;
print_ans ans_sym2;;
print_ans ans_sym3;;