open Util.Vars

module Pure =
struct

type t =
  | V of int
  | A of t * t
  | L of t
  | B

let rec print ?(l=[]) =
 function
    V n -> print_name l n
  | A(t1,t2) -> "(" ^ print ~l t1 ^ " " ^ print ~l t2 ^ ")"
  | L t ->
     let name = string_of_var (List.length l) in
     "λ" ^ name ^ "." ^ print ~l:(name::l) t
  | B -> "B"

let lift m =
 let rec aux l =
  function
   | V n -> V (if n >= l then n+m else n)
   | A (t1,t2) -> A (aux l t1, aux l t2)
   | L t -> L (aux (l+1) t)
   | B -> B
 in
  aux 0

(* Reference implementation.
   Reduction machine used below *)
let subst delift_by_one what with_what =
 let rec aux l =
  function
   | A(t1,t2) -> A(aux l t1, aux l t2)
   | V n ->
       if (if what < 0 then n = what else n = what + l) then
        lift l with_what
       else
        V (if delift_by_one && n >= l then n-1 else n)
   | L t -> L (aux (l+1) t)
   | B -> B
 in
  aux 0

(* let rec whd =
 function
  | A(t1, t2) ->
     let t2 = whd t2 in
     let t1 = whd t1 in
      (match t1 with
        | L f -> whd (subst true 0 t2 f)
        | V _
        | A _ -> A(t1,t2))
  | V _
  | L _ as t -> t
*)

let unwind ?(tbl = Hashtbl.create 317) m =
 let rec unwind (e,t,s) =
  let cache_unwind m =
   try
    Hashtbl.find tbl m
   with
    Not_found ->
     let t = unwind m in
      Hashtbl.add tbl m t ;
      t in
  let s = List.map cache_unwind s in
  let rec aux l =
   function
    | A(t1,t2) -> A(aux l t1, aux l t2)
    | V n as x when n < l -> x
    | V n ->
       (try
         lift l (cache_unwind (List.nth e (n - l)))
        with Failure _ -> V n)
    | L t -> L (aux (l+1) t)
    | B -> B in
  let t = aux 0 t in
  List.fold_left (fun f a -> A(f,a)) t s
in
 unwind m


(* let rec print_machine (e,t,s) =
 "[" ^ String.concat "," (List.map print_machine e) ^
     "|" ^ print t ^ "|" ^
     String.concat "," (List.map print_machine s) ^ "]"
;; *)

  let omega = let delta = L(A(V 0, V 0)) in A(delta,delta)

  let rec is_divergent =
   function
      t when t = omega -> true
    | A(t,_) -> is_divergent t
    | L(t) -> is_divergent t
    | _ -> false

  let mwhd m =
   let rec aux g =
    function
    (* mmm -> print_endline (print_machine mmm); match mmm with *)
        m when is_divergent (unwind m) -> [], B, []
     | (e,A(t1,t2),s) ->
         let t2' = aux g (e,t2,[]) in
         let (_,t,_) = t2' in
         if t = B
          then t2'
          else aux g (e,t1,t2'::s)
     | (e,L t,x::s) -> aux g (x::e,t,s)
     | (e,V n,s) as m ->
        (try
          let e,t,s' = List.nth e n in
          aux g (e,t,s'@s)
         with Invalid_argument _ | Failure _ -> m
         )
     | (e, B, _) -> (e, B, [])
     | (e, L t, []) ->
      let t = subst true 0 (V g) t in
      (* print_endline ("." ^ string_of_int g ^ " " ^ print_machine ((e,t,[]))); *)
      let m' = aux (g-1) (e, t, []) in
      let t' = unwind m' in
      (* print_endline ("=" ^ print t');
      print_endline ("==" ^ print (lift 1 t')); *)
      let t' = subst false g (V 0) (lift 1 t') in
      (* print_endline ("===" ^ print t'); *)
      [], (if t' = B then t' else L t'), []

   in
    unwind (aux ~-2 m)
;;

let omega should_explode =
 if should_explode
  then let f t = A(t,t) in f (L (f (V 0)))
  else B
;;

let diverged = (=) B;;

(* Note: maps only variables <= freshno *)
let env_of_sigma freshno sigma =
 let rec aux n =
  if n > freshno then
   []
  else
   let e = aux (n+1) in
   (try
    e, lift (-n-1) (snd (List.find (fun (i,_) -> i = n) sigma)),[]
   with
    Not_found -> ([], V n, []) ) :: e
 in aux 0

end

module Scott =
struct

open Pure

let rec mk_n n =
 if n = 0 then L (L (A (V 1, L (V 0)))) else L (L (A (V 0, mk_n (n-1))))

let dummy = V (max_int / 2)

let mk_match t bs =
 let bs = List.sort (fun (n1,_) (n2,_) -> compare n1 n2) bs in
 let rec aux m t =
  function
     [] -> dummy
   | (n,p)::tl as l ->
      if n = m then
       A (A (t, L (lift (m+1) p)), L (aux (m+1) (V 0) tl))
      else
       A (A (t, dummy), L (aux (m+1) (V 0) l))
 in
  aux 0 t bs

end