let print_hline = Console.print_hline;;

type var = int;;

type t =
 | V of var
 | A of t * t
 | L of t
 | LM of
    t list (* body of the match *)
  * int (* explicit liftno *)
  * int (* variable which originated this match (id) *)
 | B (* bottom *)
 | P (* pacman *)
;;


type problem = {
   orig_freshno: int
 ; freshno : int
 ; div : t list
 ; conv : t list
 ; matches : (var (* variable originating this match *) * (bool (* term comes from div*) * (t (* term to discriminate *) * t (* continuation *))) list) list
 ; sigma : (var * t) list (* substitutions *)
}

let string_of_t p =
 let bound_vars = ["x"; "y"; "z"; "z1"; "z2"] in
 let rec aux level = function
 | V v -> if v >= level then "`" ^ string_of_int (v-level) else List.nth bound_vars (level - v-1)
 | A (t1,t2) -> "("^aux level t1^" "^aux level t2^")"
 | L t -> "(\\" ^ aux (level+1) (V 0) ^ ". " ^ aux (level+1) t ^ ")"
 | LM (ts,liftno,id) -> "[match orig="^aux 0 (V id)^"]"
 | B -> "BOT"
 | P -> "PAC"
 in aux 0
;;

let string_of_t p =
  let bound_vars = ["x"; "y"; "z"; "w"; "q"; "x1"; "x2"] in
  let rec string_of_term_w_pars level = function
    | V v -> if v >= level then "`" ^ string_of_int (v-level) else List.nth bound_vars (level - v-1)
    | A _ as t -> "(" ^ string_of_term_no_pars_app level t ^ ")"
    | L _ as t -> "(" ^ string_of_term_no_pars_lam level t ^ ")"
    | LM (ts,liftno,id) -> "[match orig="^ string_of_term_w_pars 0 (V id)^"]"
    | 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 = [
  "[DV] " ^ String.concat "\n     " (List.map (string_of_t p) p.div);
  "[CV] " ^ String.concat "\n     " (List.map (string_of_t p) p.conv);
  (* "" ; *)
 ] @ Util.concat_map (fun (v, lst) -> ("[<>] of "^(string_of_t p (V v))) :: List.map (fun (b,(t,c)) -> "   " ^ (if b then "*" else " ") ^ " " ^ string_of_t p t
 (* ^" -> " ^ string_of_t p c *)
 ) lst) p.matches @ [""] in
 String.concat "\n" lines
;;

let problem_fail p reason =
 print_endline "!!!!!!!!!!!!!!! FAIL !!!!!!!!!!!!!!!";
 print_endline (string_of_problem p);
 failwith reason
;;

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

let add_to_match p id entry =
 let matches = try
   let _ = List.assoc id p.matches in
   List.map (fun (id',lst as x) -> if id <> id' then x else (id, entry::lst) ) p.matches
  with
  | Not_found -> (id,[entry]) :: p.matches in
 {p with matches}
;;

let var_occurs_in v =
 let rec aux level = function
 | V v' -> v + level = v'
 | A(t1,t2) -> (aux level t1) || (aux level t2)
 | L t -> aux (level+1) t
 | LM(ts,_,_) -> List.fold_left (fun a b -> a || aux (level+1) b) false ts
 | B -> false
 | P -> false
 in aux 0
;;

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

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

let rec subst isdiv level delift sub p =
 function
 | V v -> p, if v = level + fst sub then lift level (snd sub) else V (if delift && v > level then v-1 else v)
 | L t -> let p, t = subst isdiv (level + 1) delift sub p t in p, L t
 | A (t1,t2) ->
  let p, t1 = subst isdiv level delift sub p t1 in
  let p, t2 = subst isdiv level delift sub p t2 in
  if t1 = B || t2 = B then p, B else
  if level = 0 then mk_app isdiv p t1 t2 else p, A (t1, t2)
 | LM (ts, liftno, id) ->
   let p, ts =
    let rec aux p = function
    | [] -> p, []
    | t::ts ->
     let p, ts = aux p ts in
     let p, t = subst isdiv (level+1) delift sub p t in
     p, t::ts in
    aux p ts in
   p, LM(ts, liftno, id)
 | B -> p, B
 | P -> p, P
and mk_app isdiv p t1 t2 = let t1 = if t1 = P then L P else t1 in match t1 with
 | L t1 ->
 (* FIXME BUG HERE! *)
  if is_inert t2 && (match t2 with V v -> v > p.orig_freshno | _ -> true ) && not (var_occurs_in 0 t1)
   then (if isdiv then {p with div=p.div @ [t2]} else {p with conv=p.conv @ [t2]}), lift (-1) t1
   else subst isdiv 0 true (0, t2) p t1
 | LM(ts,liftno,id) ->
  let p, t = mk_apps isdiv p t2 (List.map (lift ~-1) ts) in
  if t = B
   then p, B
   else
    let p, cont = freshvar p in
    let newvar = V cont in
    let p = add_to_match p id (isdiv,(t,newvar)) in
      p, newvar
 | B
 | _ when t2 = B -> p, B
 | t1 -> p, A (t1, t2)
and mk_apps isdiv p t =
 function
 | [] -> p, t
 | t'::ts -> let p, t = mk_app isdiv p t t' in mk_apps isdiv p t ts
 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)
   | LM (ts, liftno, id) -> LM (List.map (aux (n'+1)) ts, n + liftno, id)
   | B -> B
   | P -> P
  in aux 0
 ;;

let subst isdiv = subst isdiv 0 false;;

let mk_lambda t = L (lift 1 t) ;;

let subst_many b sub =
 let rec aux p = function
 | [] -> p, []
 | t::tms ->
  let p, t = subst b sub p t in
  let p, tms = aux p tms in
  p, t :: tms
 in aux
;;

let rec subst_in_matches sub p =
 function
 | [] -> p, []
 | (v, branches) :: ms->
  let a, b = List.split branches in
  let b, c = List.split b in
  let p, b = subst_many false sub p b in
  let p, c = subst_many false sub p c in
  let b = List.combine b c in
  let branches = List.combine a b in
  let p, ms = subst_in_matches sub p ms in
  p, (v, branches) :: ms
;;

let subst_in_problem (sub: var * t) (p: problem) =
(* print_endline ("SUBST IN PROBLEM: " ^ string_of_t p (V (fst sub)) ^ " " ^ string_of_t p (snd sub)); *)
 let div, conv = p.div, p.conv in
 let p = {p with div=[]; conv=[]} in
 let p, div' = subst_many true sub p div in
 let p, conv' = subst_many false sub p conv in
 let p, matches = subst_in_matches sub p p.matches in
 let p = {p with div=div'@p.div; conv=conv'@p.conv; matches} in
 (* print_endline ("after sub: \n" ^ string_of_problem p); *)
 {p with sigma=sub::p.sigma}
;;

(* FIXME *)
let unify_terms p t1 t2 =
 match t1 with
 | V v -> subst_in_problem (v, t2) p
 | _ -> problem_fail p "The collapse of a match came after too many steps :(";;

let rec unify p =
 let rec give_duplicates =
  let rec aux' t = function
  | [] -> [], None
  | (b',(t',c'))::ts -> if t = t' (* FIXME! eta-eq here *)then ts, Some (b',c') else (
   let ts, res = aux' t ts in (b',(t',c'))::ts, res) in
  let rec aux = function
   | [] -> [], None
   | (b,(t,c))::rest -> (
    match aux' t rest with
    | rest, None -> aux rest
    | rest, Some (b',c') -> ((if not b' then false else b),(t,c))::rest, Some (c', c)
    ) in
  function
  | [] -> [], None
  | (orig,branches) :: ms ->
   match aux branches with
   | _, None -> let ms, res = give_duplicates ms in (orig,branches) :: ms, res
   | branches', Some subst -> (orig,branches') :: ms, Some subst in
 let matches, vars_to_be_unified = give_duplicates p.matches in
 let p = {p with matches=matches} in
 match vars_to_be_unified with
  | None -> p
  | Some(t', t) ->
   (* print_endline ("> unify " ^ string_of_t p (t') ^ " with " ^ string_of_t p t); *)
   unify (unify_terms p t' t)
;;

let problem_done p =
  let all_separated = List.for_all (fun (_, lst) -> List.length lst = 1 || List.for_all (fun (_,(t,_)) -> is_var t) lst) p.matches in
  all_separated && List.exists ((=) B) p.div
;;

let free_vars p t =
 let rec aux level = function
 | V v -> if v >= level then [v] else []
 | A(t1,t2) -> (aux level t1) @ (aux level t2)
 | L t -> aux (level+1) t
 | LM(ts,_,id) -> (List.concat (List.map (aux level) ts)) @ (
  let lst = List.assoc id p.matches in
  List.concat (List.map (fun (_,(t1,t2)) -> (aux 0 t1) @ (aux 0 t2)) lst)
 )
 | B -> []
 | P -> []
 in Util.sort_uniq (aux 0 t)
;;

let visible_vars p t =
 let rec aux = function
 | V v -> [v]
 | A(t1,t2) -> (aux t1) @ (aux t2)
 | L t -> []
 | LM(ts,_,id) -> (List.concat (List.map aux ts)) @ (
  let lst = List.assoc id p.matches in
  List.concat (List.map (fun (_,(t1,t2)) -> (aux t1) @ (aux t2)) lst)
 )
 | B -> []
 | P -> []
 in Util.sort_uniq (aux t)
;;

let forget_variable var p =
 let p', div = subst_many true (var, P) p p.div in
 let p = {p' with conv=p.conv} in
 let p, conv = subst_many false (var, B) p p.conv in
 let p, matches = subst_in_matches (var, B) p p.matches in
 {p with div; conv; matches; sigma=p.sigma}
;;
let rec remove_divergent_discriminators p =
 let f = fun (b,(t,_)) -> b && (t = B || is_lambda t) in
 try
  let v,l = List.find (fun (_,lst) -> List.exists f lst) p.matches in
  let (_,(_,c)) as entry = List.find f l in
  let l = List.filter ((<>) entry) l in
  let matches = List.map (fun (v', l') -> v', if v' = v then l else l') p.matches in
  let vars = free_vars p c in
  let p = {p with matches} in
  List.fold_right forget_variable vars p
 with Not_found -> p
;;

exception Done;;

let sanity p =
 (* try to fix divergent discriminators *)
 let p = remove_divergent_discriminators p in
 (* Remove lambdas from conv TODO remove duplicates *)
 let div = List.filter (function | P | L _ -> false | _ -> true) p.div in
 let conv = List.filter (function | B | V _ | A _ -> true | _ -> false) p.conv in
 let p = {p with div;conv} in
 (* Sanity checks: *)
 if p.div = [] then problem_fail p "p.div converged";
 if List.mem B p.conv then problem_fail p "p.conv diverged";
 if not (List.for_all (fun (_, lst) -> List.for_all (fun (b,(t,_)) -> is_inert t) lst) p.matches)
  then problem_fail p "Unsolvable discrimination";

 (* unify! :( *)
 let p = unify p in
 print_endline (string_of_problem p); (* non cancellare *)
 if problem_done p then raise Done else p
;;

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

let ignore var n p =
 print_cmd "EAT" ("on " ^ string_of_t p (V var) ^ " (of:" ^ string_of_int n ^ ")");
 let rec aux m t =
  if m = 0
   then lift n t
   else L (aux (m-1) t) in
 let p, fresh = freshvar p in
 let subst = var, aux n (V fresh) in
 sanity (subst_in_problem subst p)
;;

let explode p =
 let fv1 = List.concat (List.map (visible_vars p) p.div) in
 let fv2 = List.concat (List.map (visible_vars p) p.conv) in
 let fv = List.filter (fun x -> not (List.mem x fv2)) fv1 in
 let fv = List.filter ((<) p.orig_freshno) fv in
 match fv with
 | var::_ ->
  print_cmd "EXPLODE" ("on " ^ string_of_t p (V var));
  let subst = var, B in
  sanity (subst_in_problem subst p)
 | _ -> problem_fail p "premature explosion"
;;

let step var p =
 print_cmd "STEP" ("on " ^ string_of_t p (V var));
 let subst = var, LM([], 0, var) in
 sanity (subst_in_problem subst p)
;;

let id var p =
 print_cmd "ID" ("on " ^ string_of_t p (V var));
 let subst = var, L(V 0) in
 sanity (subst_in_problem subst p)
;;

let apply var appk p =
 print_cmd "APPLY"
  (string_of_t p (V var) ^ " applies no." ^ string_of_int appk ^ " fresh variables");
 let rec mk_freshvars n p =
  if n = 0
   then p, []
   else
    let p, vs = mk_freshvars (n-1) p in
    let p, v = freshvar p in
    p, V(v)::vs in
 let p, vars = mk_freshvars appk p in
 let p, t = mk_apps false p (V 0) (List.map (lift 1) vars) in
 let t = L (A (lift 1 (V var), t))  in
 let subst = var, t in
 sanity (subst_in_problem subst p)
;;

let perm var n p =
 print_cmd "PERM" ("on " ^ string_of_t p (V var) ^ " (of:" ^ string_of_int n ^ ")");
 let p, v = freshvar p in
 let rec aux' m t = if m < 0 then t else A(aux' (m-1) t, V m) in
 let rec aux m t =
  if m = 0
   then aux' (n-1) t
   else L (aux (m-1) t) in
 let t = aux n (lift n (V v)) in
 let subst = var, t in
 sanity (subst_in_problem subst p)
;;

(* TODO:
- cosi' come e' possibile unificare rami di branch solo se vanno -> a variabili,
  allo stesso modo e' possibile ignorare dei rami se vanno in variabili, e quelle variabili
  vengono sostituite ovunque: con bombe in conv e con pacman in div
*)
let forget var no p =
 try
  let l = List.assoc var p.matches in
  let (b,(tm,c)) = List.nth l no in
  let l = List.mapi (fun n x -> if n = no then (b,(B,c)) else x) l in
  let matches = List.map (fun (v,lst) -> v, if v = var then l else lst) p.matches in
  let p = {p with matches} in
  print_cmd "FORGET" (string_of_t p tm ^ " from the match of " ^ string_of_t p (V var));
  sanity p
  (* (match c with
   | V var ->
   print_endline (
    "--- FORGET " ^ string_of_t p tm ^ " from the match of " ^ string_of_t p (V var));
   let p = forget_variable var p in

   | _ -> print_endline "too late to forget that term"; p
  ) *)
 with Failure "nth" ->
  print_endline "wtf?"; p
;;

let parse strs =
 let dummy_p = {orig_freshno=0; freshno=0; div=[]; conv=[]; matches=[]; sigma=[]} in
  let rec aux level = function
  | Parser.Lam t -> L (aux (level + 1) t)
  | Parser.App (t1, t2) ->
   if level = 0 then snd (mk_app false dummy_p (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 rec list_split n =
 function
 | [] -> [], []
 | x::xs as l -> if n = 0 then [], l else let a, b = list_split (n-1) xs in x::a, b
;;

let magic6 div conv cmds =
 print_hline ();
 let all_tms, var_names = parse (div @ conv) in
 let div, conv = list_split (List.length div) all_tms in
 let varno = List.length var_names in
 let p = {orig_freshno=varno; freshno=1+varno; div; conv; matches=[]; sigma=[]} in
 let p = try
  let subst = Util.index_of "BOMB" var_names, L B in
  let p = subst_in_problem subst p in p
  with Not_found -> p in
 let p = sanity p in
 try
  problem_fail (List.fold_left (|>) p cmds) "Problem not completed"
 with
 | Done -> ()
;;

let interactive div conv cmds =
 print_hline ();
 let all_tms, var_names = parse (div @ conv) in
 let div, conv = list_split (List.length div) all_tms in
 let varno = List.length var_names in
 let p = {orig_freshno=varno; freshno=1+varno; div; conv; matches=[]; sigma=[]} 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 *)
 let p = sanity p in
 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 = "explode" then explode
    else if uno = "ignore" then ignore (nth spl 1) (nth spl 2)
    else if uno = "step" then step (nth spl 1)
    else if uno = "perm" then perm (nth spl 1) (nth spl 2)
    else if uno = "apply" then apply (nth spl 1) (nth spl 2)
    else if uno = "forget" then forget (nth spl 1) (nth spl 2)
    else if uno = "id" then id (nth spl 1)
    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 []
;;

let _ = magic6
 ["x x"]
 [ "_. BOMB" ]
 [ ignore 1 1; explode ]
;;

 let _ = magic6
  ["x y BOMB b"]
  [ "x BOMB y c" ]
  [ perm 1 3; step 8 ; explode; ];;

let _ = magic6
 ["x BOMB a1 c"]
 [ "x y BOMB d"; "x BOMB a2 c" ]
 [ perm 1 3 ; step 10 ; step 13; explode; ]
;;

let _ = magic6
 ["x (x x)"]
 [ "x x" ; "x x x" ] [
 apply 1 1;
 step 1;
 explode;
 step 9;
];;

let _ = magic6
 ["x (_.BOMB)"]
 [ "x (_._. BOMB)" ]
 [ apply 1 2; ]
;;

let _ = magic6
 ["x (_.y)"]
 [ "y (_. x)" ]
 [ apply 1 1; ignore 1 1;  explode; ]
;;

let _ = magic6
 ["y (x a1 BOMB c) (x BOMB b1 d)"]
 [ "y (x a2 BOMB c) (x BOMB b1 d)";
 "y (x a1 BOMB c) (x BOMB b2 d)";] [
 perm 2 3;
 step 12;
 perm 17 2;
 step 19;
 step 18;
 ignore 22 1;
 ignore 21 1;
 ignore 24 1;
 ignore 25 1;
 step 1;
 step 32;
 explode;
 ]
;;

let _ = magic6
 ["z (y x)"]
 [ "z (y (x.x))"; "y (_. BOMB)" ]
 [ apply 2 1; step 3; explode; ]
;;

let _ = magic6
 ["y x"]
 [ "y (x.x)"; "x (_. BOMB)" ]
 [ apply 1 1; ignore 2 1; step 1; explode; ]
;;

let _ = magic6
 ["z (y x)"]
 [ "z (y (x.x))"; "y (_. BOMB)" ]
 [ step 1; explode; apply 2 1; id 2; ignore 3 1; ]
;;

let _ = magic6
 ["y (x a)"]
 [ "y (x b)"; "x BOMB" ] [
 id 2;
 step 1;
 explode;
];;

magic6
 ["y (x a)"] [ "y (x b)"; "x BOMB"; "y a" ]
 [
 apply 1 1;
 perm 2 2;
 ignore 9 1;
 step 10;
 explode;
 ];;
(* "y (a c)"
[ "y (b c)"; "y (x a)"; "y (x b)"; "x BOMB" ]  *)

magic6
["x a (x (a.y BOMB))"]
[ "x b"; "x (y c)"; "x (y d)" ]
[apply 1 1;
apply 2 1;
explode;]
(* [
step 1;
step 3;
explode' 10;
(* ma si puo' fare anche senza forget *) *)
(* ] *)
;;

(* dipendente dalla codifica? no, ma si risolve solo con id *)
magic6
 ["y a"] ["y b"; "x (y (_.BOMB))"]
[
apply 1 1;
apply 2 1;
explode;
];;
 (* [id 1; explode];; *)

magic6
 ["PACMAN (x x x)"] ["PACMAN (x x)"]
 [
 ignore 2 2;
 explode;
 ];;

print_hline();
print_endline "ALL DONE. "