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_flag = bool ;;

type var = int;;
type t =
 | V of var
 | A of var_flag * t * t
 | L of t
;;

let string_of_t =
  let sep_of_app b = if b then " +" else " " in
  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 "v" ^ (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)
    | A _
    | L _ as t -> "(" ^ string_of_term_no_pars level t ^ ")"
  and string_of_term_no_pars_app level = function
    | A(b,t1,t2) -> string_of_term_no_pars_app level t1 ^ sep_of_app b ^ 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
;;

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

let measure_of_t =
 let rec aux = function
 | V _ -> 0
 | A(b,t1,t2) ->
   (if b then 1 else 0) + aux t1 + aux t2
 | L t -> aux t
 in aux
;;

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

let string_of_problem p =
 let measure = List.fold_left (+) 0 (List.map measure_of_t p.tms) in
 let lines = ("[measure] " ^ string_of_int measure) ::
   List.map (fun x -> "[TM] " ^ string_of_t x) p.tms 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
 | L _ -> 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
;;

(* precomputes the number of leading lambdas in a term,
   after replacing _v_ w/ a term starting with n lambdas *)
let rec no_leading_lambdas v n = function
 | L t -> 1 + no_leading_lambdas (v+1) n t
 | A _ as t -> let v', m = get_inert t in if v = v' then max 0 (n - m) else 0
 | V v' -> if v = v' then n else 0
;;

let rec erase = function
 | L t -> L (erase t)
 | A(_,t1,t2) -> A(false, erase t1, erase t2)
 | V _ as t -> t
;;

let explode =
 let rec aux args = function
 | L _ -> assert false
 | V _ as x -> x, args
 | A(b,t1,t2) -> aux ((b,t2)::args) t1
 in aux []
;;

let rec implode hd args =
 match args with
 | [] -> hd
 | (f,a)::args -> implode (A(f,hd,a)) args
;;

let get_head =
 let rec aux lev = function
 | L t -> aux (lev+1) t
 | A(_,t,_) -> aux lev t
 | V v -> v - lev
 in aux 0
;;

let rec subst level delift ((var, tm) as sub) =
 function
 | V v -> if v = level + var then lift level tm else V (if delift && v > level then v-1 else v)
 | L t -> L (subst (level + 1) delift sub t)
 | A(b,t1,t2) ->
    let t1' = subst level delift sub t1 in
    let t2' = subst level delift sub t2 in
    mk_app b t1' t2'
and mk_app flag t1 t2 = match t1 with
 | L t1 -> subst 0 true (0, t2) t1
 | _ -> A (flag, 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 (b,t1, t2) -> A (b,aux lev t1, aux lev t2)
 in aux 0
;;
let subst = subst 0 false;;
(* let mk_app = mk_app true;; *)
let rec mk_apps t = function
 | [] -> t
 | (f,x)::xs -> mk_apps (mk_app f t x) xs
;;

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

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

let subst_in_problem ?(top=true) ((v, t) as sub) p =
print_endline ("-- SUBST " ^ string_of_t (V v) ^ " |-> " ^ string_of_t t);
 let sigma = sub::p.sigma in
 let sub = (v, t) in
 let tms = List.map (subst sub) p.tms in
 {p with tms; sigma}
;;

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

let check p sigma =
 assert false (* FIXME *)
 (* print_endline "Checking...";
 let tms = List.map purify p.tms in
 let sigma = List.map (fun (v,t) -> v, purify t) sigma in
 let freshno = List.fold_right (max ++ fst) 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 *)
 let rec all_different = function
  | [] -> true
  | x::xs -> List.for_all ((<>) x) xs && all_different xs in
 if List.for_all is_var p.tms && all_different p.tms
  then raise (Done p.sigma);
 if List.exists (not ++ is_inert) p.tms
  then problem_fail p "used a non-effective path";
 p
;;

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

let step var j n p =
  let atsnd f (a,b) = (a, f b) in
  let p, alphas = (* make fresh vars *)
   fold_nat (fun (p, vs) _ ->
     let p, v = freshvar p in
     p, v::vs
   ) (p, []) n in let alphas = List.rev alphas in
  let rec aux lev (inside:bool) = function
  | L t -> L (aux (lev+1) inside t)
  | _ as x ->
    let hd, args = explode x in
    if hd = V (var+lev) then
      (let nargs = List.length args in
       let k = max 0 (j + 1 - nargs) in
       let args = List.mapi
        (fun i (f, t) -> f, lift k (aux lev (if i=j then true else inside) t)) args in
       let bound = fold_nat (fun x n -> (false,V(n-1)) :: x) [] k in
       let args = args @ bound in
       let _, head = List.nth args j in
       let args = List.mapi
        (fun i (f, t) -> (if i=j && not inside then false else f), if i=j && not inside then erase t else t) args in
       let head = (if inside then erase else id) head in
       print_endline ("HEAD: " ^ string_of_t head);
       let alphas = List.map (fun v -> false, V(lev+k+v)) alphas in
       let t = mk_apps head (alphas @ args) in
       let t = mk_lams t k in
       t
    ) else
    (let args = List.map (atsnd (aux lev inside)) args in
     implode hd args) in
  let sigma = (var, aux 0 false (V var)) :: p.sigma in
  {p with tms=List.map (aux 0 false) p.tms; sigma}
;;

let finish p = assert false ;;

let rec auto p = assert false ;;

let problem_of (label, div, convs, ps, var_names) =
 print_hline ();
 let rec aux = function
 | `Lam(_,t) -> L (aux t)
 | `I ((v,_), args) -> Listx.fold_left (fun x y -> mk_app true x (aux y)) (V v) args
 | `Var(v,_) -> V v
 | `N _ | `Match _ -> assert false in
 let convs = (List.rev convs :> Num.nf list) in
 let tms = List.map aux (convs @ (ps :> Num.nf list)) in
 let tms = match div with
 | Some div -> aux (div :> Num.nf) :: tms
 | None -> tms in
 let varno = List.length var_names in
 let p = {orig_freshno=varno; freshno=1+varno; tms; sigma=[]} in
 (* initial sanity check *)
 sanity p
;;

let rec interactive p =
 print_string "[varno index alphano] ";
 let s = read_line () in
 let spl = Str.split (Str.regexp " +") s in
 let nth n = int_of_string (List.nth spl n) in
 let p = step (nth 0) (nth 1) (nth 2) p in
 interactive (sanity p)
;;

let solve p =
 let rec aux = function
 | [] -> false
 | x::xs -> List.exists (eta_subterm x) xs || aux xs in
 if aux p.tms
  then print_endline "!!! Problem stopped: subterm problem !!!"
  else check p (interactive p)
;;

Problems.main (solve ++ problem_of);