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 = {
   freshno : int
 ; div : t
 ; conv : t list
 ; matches : (var (* variable originating this match *) * (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_problem p =
 let lines = [
  "[DV] " ^ string_of_t p p.div;
  "[CV] " ^ String.concat "\n     " (List.map (string_of_t p) p.conv);
  "" ;
 ] @ Util.concat_map (fun (_, lst) -> "[<>]" :: List.map (fun (t,c) -> "     " ^ 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 FAIL FAIL FAIL FAIL FAIL FAIL FAIL 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=matches}
;;

let free_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
 | L _ | LM _ | B -> false
 | _ -> true
;;

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

let rec subst 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 (level + 1) delift sub p t in p, L t
 | A (t1,t2) ->
  let p, t1 = subst level delift sub p t1 in
  let p, t2 = subst level delift sub p t2 in
  if t1 = B || t2 = B then p, B else
  if level = 0 then mk_app 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 (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 p t1 t2 = let t1 = if t1 = P then L P else t1 in match t1 with
 | L t1 ->
  let p = if is_inert t2 && not (is_var t2) && free_in 0 t1 then {p with conv=t2::p.conv} else p in
   subst 0 true (0, t2) p t1
 | LM(ts,liftno,id) ->
  let p, t = mk_apps 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 (t,newvar) in
      p, newvar
 | B
 | _ when t2 = B -> p, B
 | t1 -> p, A (t1, t2)
and mk_apps p t =
 function
 | [] -> p, t
 | t'::ts -> let p, t = mk_app p t t' in mk_apps p t ts
and lift n = function
 | V m -> V (m + n)
 | L t -> L (lift n t)
 | A (t1, t2) -> A (lift n t1, lift n t2)
 | LM (ts, liftno, id) -> LM (List.map (lift n) ts, n + liftno, id)
 | B -> B
 | P -> P
 ;;

let subst = subst 0 false;;

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

let subst_many sub =
 let rec aux p = function
 | [] -> p, []
 | t::tms ->
  let p, t = subst 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 p, a = subst_many sub p a in
  let p, b = subst_many sub p b 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 p', div = subst sub p p.div in
 let p = {p' with conv=p.conv} in
 let p, conv = subst_many sub p p.conv in
 let p, matches = subst_in_matches sub p p.matches in
 {p with div=div; conv=conv; matches=matches; 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
  | (t',c')::ts -> if t = t' then ts, Some c' else (
   let ts, res = aux' t ts in (t',c')::ts, res) in
  let rec aux = function
   | [] -> [], None
   | (t,c)::rest -> (
    match aux' t rest with
    | rest, None -> aux rest
    | rest, Some c' -> (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 && p.div = B
;;

exception Done;;

let sanity p =
 (* Sanity checks: *)
 if is_lambda p.div || is_pacman p.div then problem_fail p "p.div converged" ;
 List.iter (function | B -> problem_fail p "p.conv diverged" | _ -> ()) p.conv;
 if not (List.for_all (fun (_, lst) -> List.for_all (fun (t,_) -> is_inert t) lst) p.matches)
  then problem_fail p "Lambda in a match: k was not big enough?";

 (* Remove lambdas from conv TODO remove duplicates *)
 let conv = List.filter is_inert p.conv in
 let p = {p with conv=conv} in
 (* unify! :( *)
 let p = unify p in
 print_endline (string_of_problem p);
 if problem_done p then raise Done else p
;;

let ignore var n p =
print_endline (
 "--- 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 subst = var, aux n (V var) in
 sanity (subst_in_problem subst p)
;;

let explode' var p =
 print_endline (
 "--- EXPLODE on " ^ string_of_t p (V var)
 );
 let subst = var, B in
 sanity (subst_in_problem subst p)
 ;;

let explode p =
 match p.div with
 | V var -> explode' var p
 | _ -> problem_fail p "premature explosion"
;;

let step var p =
 print_endline (
  "--- 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_endline (
  "--- 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_endline (
  "--- APPLY_CONSTS on " ^ string_of_t p (V var)
  ^ " (special_k:" ^ string_of_int appk ^ ")");
 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 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_endline (
  "--- 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 var)) 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 n p =
 let remove_from_branch p var n = ... in
 let p, (tm, c) = remove_from_branch p var n in
 print_endline (
  "--- FORGET " ^ string_of_t p tm ^ " from the match of " ^ string_of_t p (V var)
  ^ " (of:" ^ string_of_int n ^ ")");
 sanity p
;; *)

let parse strs =
 let dummy_p = {freshno=0; div=V 0; 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 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 magic6 div conv cmds =
 print_hline ();
 let all_tms, var_names = parse (div :: conv) in
 let div, conv = List.hd all_tms, List.tl all_tms in
 let freshno = 1 + List.length var_names in
 let p = {freshno; 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.hd all_tms, List.tl all_tms in
 let freshno = 1 + List.length var_names in
 let p = {freshno; 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
 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 = "id" then id (nth spl 1)
    else failwith "unknown"
    with Failure _ -> print_endline "Wrong input."; (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 -> 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 1 ; ignore 8 2; explode; ];;

let _ = magic6
   "x BOMB a1 c"
 [ "x y BOMB d"; "x BOMB a2 c" ]
 [ perm 1 3 ; step 1 ; step 12; ignore 13 1; 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 2; step 17; perm 16 2; step 16; ignore 19 1; ignore 20 1; ignore 22 1; ignore 23 1; step 1; step 26; explode; ]
;;

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

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;
 perm 2 2;
 step 3;
 apply 2 1;
 ignore 4 1;
 step 2;
 ignore 12 1;
 ignore 13 1;
 step 1;
 explode;
];;

interactive
  "y (x a)"
[ "y (x b)"; "x BOMB"; "y a" ] [];;

print_endline "ALL DONE. "