let (++) f g x = f (g x);;
let id x = x;;

let print_hline = Console.print_hline;;

type var = int;;
type t =
 | V of var
 | A of t * t
 | L of t
 | B (* bottom *)
 | P (* pacman *)
;;

let eta_eq =
 let rec aux l1 l2 t1 t2 = match t1, t2 with
  | L t1, L t2 -> aux l1 l2 t1 t2
  | L t1, t2 -> aux l1 (l2+1) t1 t2
  | t1, L t2 -> aux (l1+1) l2 t1 t2
  | V a, V b -> a + l1 = b + l2
  | A(t1,t2), A(u1,u2) -> aux l1 l2 t1 u1 && aux l1 l2 t2 u2
  | _, _ -> false
 in aux 0 0
;;

type problem = {
   orig_freshno: int
 ; freshno : int
 ; div : t
 ; conv : t
 ; sigma : (var * t) list (* substitutions *)
 ; stepped : var list
}

exception Done of (var * t) list (* substitution *);;
exception Fail of int * string;;

let string_of_t p =
  let bound_vars = ["x"; "y"; "z"; "w"; "q"] in
  let rec string_of_term_w_pars level = function
    | V v -> if v >= level then "`" ^ string_of_int (v-level) else
     let nn = level - v-1 in
      if nn < 5 then List.nth bound_vars nn else "x" ^ (string_of_int (nn-4))
    | A _
    | L _ as t -> "(" ^ string_of_term_no_pars_lam level t ^ ")"
    | B -> "BOT"
    | P -> "PAC"
  and string_of_term_no_pars_app level = function
    | A(t1,t2) -> (string_of_term_no_pars_app level t1) ^ " " ^ (string_of_term_w_pars level t2)
    | _ as t -> string_of_term_w_pars level t
  and string_of_term_no_pars_lam level = function
    | L t -> "λ" ^ string_of_term_w_pars (level+1) (V 0) ^ ". " ^ (string_of_term_no_pars_lam (level+1) t)
    | _ as t -> string_of_term_no_pars level t
  and string_of_term_no_pars level = function
    | L _ as t -> string_of_term_no_pars_lam level t
    | _ as t -> string_of_term_no_pars_app level t
  in string_of_term_no_pars 0
;;

let string_of_problem p =
 let lines = [
  "[stepped] " ^ String.concat " " (List.map string_of_int p.stepped);
  "[DV] " ^ (string_of_t p p.div);
  "[CV] " ^ (string_of_t p p.conv);
 ] in
 String.concat "\n" lines
;;

let problem_fail p reason =
 print_endline "!!!!!!!!!!!!!!! FAIL !!!!!!!!!!!!!!!";
 print_endline (string_of_problem p);
 raise (Fail (-1, reason))
;;

let freshvar ({freshno} as p) =
 {p with freshno=freshno+1}, freshno+1
;;

let rec is_inert =
 function
 | A(t,_) -> is_inert t
 | V _ -> true
 | L _ | B | P -> false
;;

let is_var = function V _ -> true | _ -> false;;
let is_lambda = function L _ -> true | _ -> false;;

let rec head_of_inert = function
 | V n -> n
 | A(t, _) -> head_of_inert t
 | _ -> assert false
;;

let rec args_no = function
 | V _ -> 0
 | A(t, _) -> 1 + args_no t
 | _ -> assert false
;;

let rec subst level delift fromdiv sub =
 function
 | V v -> if v = level + fst sub then lift level (snd sub) else V (if delift && v > level then v-1 else v)
 | L t -> L (subst (level + 1) delift fromdiv sub t)
 | A (t1,t2) ->
  let t1 = subst level delift fromdiv sub t1 in
  let t2 = subst level delift fromdiv sub t2 in
  if t1 = B || t2 = B then B else mk_app fromdiv t1 t2
 | B -> B
 | P -> P
and mk_app fromdiv t1 t2 = let t1 = if t1 = P then L P else t1 in match t1 with
 | B | _ when t2 = B -> B
 | L t1 -> subst 0 true fromdiv (0, t2) t1
 | t1 -> A (t1, t2)
and lift n =
 let rec aux n' =
  function
  | V m -> V (if m >= n' then m + n else m)
  | L t -> L (aux (n'+1) t)
  | A (t1, t2) -> A (aux n' t1, aux n' t2)
  | B -> B
  | P -> P
 in aux 0
;;
let subst = subst 0 false;;

let subst_in_problem (sub: var * t) (p: problem) =
print_endline ("-- SUBST " ^ string_of_t p (V (fst sub)) ^ " |-> " ^ string_of_t p (snd sub));
 let p = {p with stepped=(fst sub)::p.stepped} in
 let conv = subst false sub p.conv in
 let div = subst true sub p.div in
 let p = {p with div; conv} in
 (* print_endline ("after sub: \n" ^ string_of_problem p); *)
 {p with sigma=sub::p.sigma}
;;

let get_subterm_with_head_and_args hd_var n_args =
 let rec aux = function
 | V _ | L _ | B | P -> None
 | A(t1,t2) as t ->
   if head_of_inert t1 = hd_var && n_args <= 1 + args_no t1
    then Some t
    else match aux t2 with
    | None -> aux t1
    | Some _ as res -> res
 in aux
;;

(* let rec simple_explode p =
 match p.div with
 | V var ->
  let subst = var, B in
  sanity (subst_in_problem subst p)
 | _ -> p *)

let sanity p =
 print_endline (string_of_problem p); (* non cancellare *)
 if p.div = B then raise (Done p.sigma);
 if not (is_inert p.div) then problem_fail p "p.div converged";
 if p.conv = B then problem_fail p "p.conv diverged";
 (* let p = if is_var p.div then simple_explode p else p in *)
 p
;;

let print_cmd s1 s2 = print_endline (">> " ^ s1 ^ " " ^ s2);;

(* eat the arguments of the divergent and explode.
 It does NOT perform any check, may fail if done unsafely *)
let eat p =
print_cmd "EAT" "";
 let var = head_of_inert p.div in
 let n = args_no p.div in
 let rec aux m t =
  if m = 0
   then lift n t
   else L (aux (m-1) t) in
 let subst = var, aux n B in
 sanity (subst_in_problem subst p)
;;

(* step on the head of div, on the k-th argument, with n fresh vars *)
let step k n p =
 let var = head_of_inert p.div in
 print_cmd "STEP" ("on " ^ string_of_t p (V var) ^ " (of:" ^ string_of_int n ^ ")");
 let rec aux' p m t =
  if m < 0
   then p, t
    else
     let p, v = freshvar p in
     let p, t = aux' p (m-1) t in
      p, A(t, V (v + k + 1)) in
 let p, t = aux' p n (V 0) in
 let rec aux' m t = if m < 0 then t else A(aux' (m-1) t, V (k-m)) in
 let rec aux m t =
  if m < 0
   then aux' (k-1) t
   else L (aux (m-1) t) in
 let t = aux k t in
 let subst = var, t in
 sanity (subst_in_problem subst p)
;;

let parse strs =
  let rec aux level = function
  | Parser.Lam t -> L (aux (level + 1) t)
  | Parser.App (t1, t2) ->
   if level = 0 then mk_app false (aux level t1) (aux level t2)
    else A(aux level t1, aux level t2)
  | Parser.Var v -> V v
  in let (tms, free) = Parser.parse_many strs
  in (List.map (aux 0) tms, free)
;;

let problem_of div conv =
 print_hline ();
 let all_tms, var_names = parse ([div; conv]) in
 let div, conv = List.hd all_tms, List.hd (List.tl all_tms) in
 let varno = List.length var_names in
 let p = {orig_freshno=varno; freshno=1+varno; div; conv; sigma=[]; stepped=[]} in
 (* activate bombs *)
 let p = try
  let subst = Util.index_of "BOMB" var_names, L B in
   subst_in_problem subst p
  with Not_found -> p in
 (* activate pacmans *)
 let p = try
  let subst = Util.index_of "PACMAN" var_names, P in
   let p = subst_in_problem subst p in
   (print_endline ("after subst in problem " ^ string_of_problem p); p)
  with Not_found -> p in
 (* initial sanity check *)
 sanity p
;;

let exec div conv cmds =
 let p = problem_of div conv in
 try
  problem_fail (List.fold_left (|>) p cmds) "Problem not completed"
 with
 | Done _ -> ()
;;

let inert_cut_at n t =
 let rec aux t =
  match t with
  | V _ as t -> 0, t
  | A(t1,_) as t ->
    let k', t' = aux t1 in
     if k' = n then n, t'
      else k'+1, t
  | _ -> assert false
 in snd (aux t)
;;

let find_eta_difference p t n_args =
 let t = inert_cut_at n_args t in
 let rec aux t u k = match t, u with
 | V _, V _ -> assert false (* div subterm of conv *)
 | A(t1,t2), A(u1,u2) ->
    if not (eta_eq t2 u2) then (print_endline((string_of_t p t2) ^ " <> " ^ (string_of_t p u2)); k)
    else aux t1 u1 (k-1)
 | _, _ -> assert false
 in aux p.div t n_args
;;

let rec no_leading_lambdas = function
 | L t -> 1 + no_leading_lambdas t
 | _ -> 0
;;

let compute_max_lambdas_at hd_var j =
 let rec aux hd = function
 | A(t1,t2) ->
    (if head_of_inert t1 = hd && args_no t1 = j
      then max (
       if is_inert t2 && head_of_inert t2 = hd
        then j - args_no t2
        else no_leading_lambdas t2)
      else id) (max (aux hd t1) (aux hd t2))
 | L t -> aux (hd+1) t
 | V _ -> 0
 | _ -> assert false
 in aux hd_var
;;

let rec auto p =
 let hd_var = head_of_inert p.div in
 let n_args = args_no p.div in
 match get_subterm_with_head_and_args hd_var n_args p.conv with
 | None ->
    (try let p = eat p in problem_fail p "Auto did not complete the problem"  with Done _ -> ())
 | Some t ->
  let j = find_eta_difference p t n_args - 1 in
  let k = max
   (compute_max_lambdas_at hd_var j p.div)
    (compute_max_lambdas_at hd_var j p.conv) in
  let p = step j k p in
  auto p
;;

let interactive div conv cmds =
 let p = problem_of div conv in
 try (
 let p = List.fold_left (|>) p cmds in
 let rec f p cmds =
  let nth spl n = int_of_string (List.nth spl n) in
  let read_cmd () =
   let s = read_line () in
   let spl = Str.split (Str.regexp " +") s in
   s, let uno = List.hd spl in
    try if uno = "eat" then eat
    else if uno = "step" then step (nth spl 1) (nth spl 2)
    else failwith "Wrong input."
    with Failure s -> print_endline s; (fun x -> x) in
  let str, cmd = read_cmd () in
  let cmds = (" " ^ str ^ ";")::cmds in
  try
   let p = cmd p in f p cmds
  with
  | Done _ -> print_endline "Done! Commands history: "; List.iter print_endline (List.rev cmds)
 in f p []
 ) with Done _ -> ()
;;

let rec conv_join = function
 | [] -> "@"
 | x::xs -> conv_join xs ^ " ("^ x ^")"
;;

let _ = exec
 "x x"
 (conv_join["x y"; "y y"; "y x"])
 [ step 0 0; eat ]
;;

auto (problem_of "x x" "@ (x y) (y y) (y x)");;
auto (problem_of "x y" "@ (x (_. x)) (y z) (y x)");;
auto (problem_of "a (x. x b) (x. x c)" "@ (a (x. b b) @) (a @ c) (a (x. x x) a) (a (a a a) (a c c))");;

interactive "x y"
"@ (x x) (y x) (y z)" [step 0 0; step 0 1; eat] ;;

auto (problem_of "x (y. x y y)" "x (y. x y x)");;

auto (problem_of "x a a a a" (conv_join[
 "x b a a a";
 "x a b a a";
 "x a a b a";
 "x a a a b";
]));

(* Controesempio ad usare un conto dei lambda che non considere le permutazioni *)
auto (problem_of "x a a a a (x (x. x x) @ @ (_._.x. x x) x) b b b" (conv_join[
 "x a a a a (_. a) b b b";
 "x a a a a (_. _. _. _. x. y. x y)";
]));


print_hline();
print_endline "ALL DONE. "

(* TEMPORARY TESTING FACILITY BELOW HERE *)

let acaso l =
    let n = Random.int (List.length l) in
    List.nth l n
;;

let acaso2 l1 l2 =
  let n1 = List.length l1 in
  let n = Random.int (n1 + List.length l2) in
  if n >= n1 then List.nth l2 (n - n1) else List.nth l1 n
;;

let gen n vars =
  let rec aux n inerts lams =
    if n = 0 then List.hd inerts, List.hd (Util.sort_uniq (List.tl inerts))
    else let inerts, lams = if Random.int 2 = 0
      then inerts, ("(" ^ acaso vars ^ ". " ^ acaso2 inerts lams ^ ")") :: lams
      else ("(" ^ acaso inerts ^ " " ^ acaso2 inerts lams^ ")") :: inerts, lams
    in aux (n-1) inerts lams
  in aux (2*n) vars []
;;

let f () =
  let complex = 200 in
  let vars = ["x"; "y"; "z"; "v" ; "w"; "a"; "b"; "c"] in
  gen complex vars

  let rec repeat f n =
    prerr_endline "\n########################### NEW TEST ###########################";
    f () ;
    if n > 0 then repeat f (n-1)
  ;;


let main () =
 Random.self_init ();
 repeat (fun _ ->
   let div, conv = f () in
   auto (problem_of div conv)
 ) 100;
;;

(* main ();; *)