let (++) f g x = f (g x);;
let id x = x;;
let rec fold_nat f x n = if n = 0 then x else f (fold_nat f x (n-1)) n ;;

let print_hline = Console.print_hline;;

open Pure

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

let delta = L(A(V 0, V 0));;

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
  | C a, C b -> a = b
  | A(t1,t2), A(u1,u2) -> aux l1 l2 t1 u1 && aux l1 l2 t2 u2
  | _, _ -> false
 in aux ;;
let eta_eq = eta_eq' 0 0;;

(* is arg1 eta-subterm of arg2 ? *)
let eta_subterm u =
 let rec aux lev t = eta_eq' lev 0 u t || match t with
 | L t -> aux (lev+1) t
 | A(t1, t2) -> aux lev t1 || aux lev t2
 | _ -> false
 in aux 0
;;

(* does NOT lift the argument *)
let mk_lams = fold_nat (fun x _ -> L x) ;;

let string_of_t =
  let string_of_bvar =
   let bound_vars = ["x"; "y"; "z"; "w"; "q"] in
   let bvarsno = List.length bound_vars in
   fun nn -> if nn < bvarsno then List.nth bound_vars nn else "x" ^ (string_of_int (nn - bvarsno + 1)) in
  let rec string_of_term_w_pars level = function
    | V v -> if v >= level then "`" ^ string_of_int (v-level) else
       string_of_bvar (level - v-1)
    | C n -> "c" ^ string_of_int n
    | A _
    | L _ as t -> "(" ^ string_of_term_no_pars level t ^ ")"
    | B -> "BOT"
  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 level = function
    | L t -> "λ" ^ string_of_bvar level ^ ". " ^ string_of_term_no_pars (level+1) t
    | _ as t -> string_of_term_no_pars_app level t
  in string_of_term_no_pars 0
;;

type problem = {
   orig_freshno: int
 ; freshno : int
 ; div : t
 ; conv : t
 ; sigma : (var * t) list (* substitutions *)
 ; stepped : var list
 ; phase : [`One | `Two] (* :'( *)
}

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

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

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
 | C _
 | L _ | B -> false
;;

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

let rec get_inert = function
 | V n -> (n,0)
 | A(t, _) -> let hd,args = get_inert t in hd,args+1
 | _ -> assert false
;;

let rec no_leading_lambdas hd_var j = function
 | L t -> 1 + no_leading_lambdas (hd_var+1) j t
 | A _ as t -> let hd_var', n = get_inert t in if hd_var = hd_var' then max 0 (j - n) else 0
 | V n -> if n = hd_var then j else 0
 | B | C _ -> 0
;;
let rec subst level delift 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 -> let t = subst (level + 1) delift sub t in if t = B then B else L t
 | A (t1,t2) ->
  let t1 = subst level delift sub t1 in
  let t2 = subst level delift sub t2 in
  mk_app t1 t2
 | C _ as t -> t
 | B -> B
and mk_app t1 t2 = if t2 = B || (t1 = delta && t2 = delta) then B
 else match t1 with
 | B -> B
 | L t1 -> subst 0 true (0, t2) t1
 | _ -> A (t1, t2)
and lift n =
 let rec aux lev =
  function
  | V m -> V (if m >= lev then m + n else m)
  | L t -> L (aux (lev+1) t)
  | A (t1, t2) -> A (aux lev t1, aux lev t2)
  | C _ as t -> t
  | B -> B
 in aux 0
;;
let subst = subst 0 false;;

let subst_in_problem sub p =
print_endline ("-- SUBST " ^ string_of_t (V (fst sub)) ^ " |-> " ^ string_of_t (snd sub));
 {p with
  div=subst sub p.div;
  conv=subst sub p.conv;
  stepped=(fst sub)::p.stepped;
  sigma=sub::p.sigma}
;;

let get_subterm_with_head_and_args hd_var n_args =
 let rec aux lev = function
 | C _
 | V _ | B -> None
 | L t -> aux (lev+1) t
 | A(t1,t2) as t ->
   let hd_var', n_args' = get_inert t1 in
   if hd_var' = hd_var + lev && n_args <= 1 + n_args'
    then Some (lift ~-lev t)
    else match aux lev t2 with
    | None -> aux lev t1
    | Some _ as res -> res
 in aux 0
;;

let rec purify = function
 | L t -> Pure.L (purify t)
 | A (t1,t2) -> Pure.A (purify t1, purify t2)
 | V n -> Pure.V n
 | C _ -> Pure.V max_int (* FIXME *)
 | B -> Pure.B
;;

let check p sigma =
 print_endline "Checking...";
 let div = purify p.div in
 let conv = purify p.conv in
 let sigma = List.map (fun (v,t) -> v, purify t) sigma in
 let freshno = List.fold_right (fun (x,_) -> max x) sigma 0 in
 let env = Pure.env_of_sigma freshno sigma in
 assert (Pure.diverged (Pure.mwhd (env,div,[])));
 print_endline " D diverged.";
 assert (not (Pure.diverged (Pure.mwhd (env,conv,[]))));
 print_endline " C converged.";
 ()
;;

let sanity p =
 print_endline (string_of_problem p); (* non cancellare *)
 if p.conv = B then problem_fail p "p.conv diverged";
 if p.div = B then raise (Done p.sigma);
 if p.phase = `Two && p.div = delta then raise (Done p.sigma);
 if not (is_inert p.div) then problem_fail p "p.div converged";
 p
;;

(* drops the arguments of t after the n-th *)
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 t2) ^ " <> " ^ (string_of_t u2));*) k)
    else aux t1 u1 (k-1)
 | _, _ -> assert false
 in aux p.div t n_args
;;

let compute_max_lambdas_at hd_var j =
 let rec aux hd = function
 | A(t1,t2) ->
    (if get_inert t1 = (hd, j)
      then max ( (*FIXME*)
       if is_inert t2 && let hd', j' = get_inert t2 in hd' = hd
        then let hd', j' = get_inert t2 in j - j'
        else no_leading_lambdas hd_var j 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 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, k = get_inert p.div in
 let phase = p.phase in
 let p, t =
  match phase with
  | `One ->
      let n = 1 + max
       (compute_max_lambdas_at var k p.div)
       (compute_max_lambdas_at var k p.conv) in
      (* apply fresh vars *)
      let p, t = fold_nat (fun (p, t) _ ->
        let p, v = freshvar p in
        p, A(t, V (v + k))
      ) (p, V 0) n in
      let p = {p with phase=`Two} in p, A(t, delta)
  | `Two -> p, delta in
 let subst = var, mk_lams t k in
 let p = subst_in_problem subst p in
 sanity p;
 let p = if phase = `One then {p with div = (match p.div with A(t,_) -> t | _ -> assert false)} else p in
 sanity p
;;

(* step on the head of div, on the k-th argument, with n fresh vars *)
let step k n p =
 let var, _ = get_inert p.div in
print_cmd "STEP" ("on " ^ string_of_t (V var) ^ " (of:" ^ string_of_int n ^ ")");
 let p, t = (* apply fresh vars *)
  fold_nat (fun (p, t) _ ->
    let p, v = freshvar p in
    p, A(t, V (v + k + 1))
  ) (p, V 0) n in
 let t = (* apply unused bound variables V_{k-1}..V_1 *)
  fold_nat (fun t m -> A(t, V (k-m+1))) t k in
 let t = mk_lams t (k+1) in (* make leading lambdas *)
 let subst = var, t in
 let p = subst_in_problem subst p in
 sanity p
;;
;;

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

let rec auto p =
 let hd_var, n_args = get_inert p.div in
 match get_subterm_with_head_and_args hd_var n_args p.conv with
 | None ->
   (try
    let phase = p.phase in
     let p = eat p in
     if phase = `Two
      then problem_fail p "Auto.2 did not complete the problem"
      else auto p
    with Done sigma -> sigma)
 | Some t ->
  let j = find_eta_difference p t n_args - 1 in
  let k = 1 + 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 problem_of div convs =
 let rec conv_join = function
  | [] -> "@"
  | x::xs -> conv_join xs ^ " ("^ x ^")" in
 print_hline ();
 let conv = conv_join convs in
 let [@warning "-8"] [div; conv], var_names = parse ([div; conv]) in
 let varno = List.length var_names in
 let p = {orig_freshno=varno; freshno=1+varno; div; conv; sigma=[]; stepped=[]; phase=`One} in
 (* initial sanity check *)
 sanity p
;;

let solve p =
 if eta_subterm p.div p.conv
  then print_endline "!!! div is subterm of conv. Problem was not run !!!"
  else check p (auto p)
;;

let run x y = solve (problem_of x y);;

(* Example usage of interactive: *)

(* 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 _ -> ()
;; *)

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

run "x x" ["x y"; "y y"; "y x"] ;;
run "x y" ["x (_. x)"; "y z"; "y x"] ;;
run "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)"] ;;

run "x (y. x y y)" ["x (y. x y x)"] ;;

run "x a a a a" [
 "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 *)
run "x a a a a (x (x. x x) @ @ (_._.x. x x) x) b b b" [
 "x a a a a (_. a) b b b";
 "x a a a a (_. _. _. _. x. y. x y)";
] ;;

(* bug in no_leading_lambdas where x in j-th position *)
run "v0 (v1 v2) (x. y. v0 v3 v4 v2)"
["v5 (v6 (x. v6) v7 v8 (x. v9) (x. v9 (v10 v10))) (v6 (x. v6) v7 v8 (v9 (x. v11) (v4 v8) (x. y. z. v12)) (v13 (x. y. y))) (v9 (x. v11) (v4 v8) (x. y. z. v12) (v4 v6 (x. v7 v11) v12) (x. x v8)) (v0 v3 v4 v9 (v7 v11) (v0 v3 v4 v2) v10) (v6 (x. v6) v7 v8 (v9 (x. v11) (v4 v8) (x. y. z. v12)) (v13 (x. y. y)) (v9 (x. v11) (v4 v8)))"]
;;

print_hline();
print_endline ">>> EXAMPLES IN simple.ml FINISHED <<<"