open Util
open Util.Vars
open Pure

(* debug options *)
let debug_display_arities = false;;

(************ Syntax ************************************)

(* Normal forms*)

(* Var n = n-th De Bruijn index, 0-based *)

(*type nf =
  | Lam of nf
  | Var of int
  | i
and i =
  | I of int * nf listx
;;*)
type var = int * (* arity of variable*) int;;
type 'nf_nob i_var_ = [ `I of var * 'nf_nob Listx.listx | `Var of var ]
type 'nf_nob i_n_var_ = [ `N of int | 'nf_nob i_var_ ]
type ('nf_nob,'nf) i_num_var_ = [
  | 'nf_nob i_n_var_
  | `Match of ('nf_nob,'nf) i_num_var_ * (* originating var *) var * (*lift*) int * (*branches*)(int * 'nf) list ref * (*args*)'nf_nob list
]
type 'nf nf_nob_ = [ `Lam of (* was_unpacked *) bool * 'nf | `Pacman | ('nf nf_nob_,'nf) i_num_var_ ]
type nf = [ nf nf_nob_ | `Bottom  ]
type nf_nob = nf nf_nob_
type i_var = nf_nob i_var_;;
type i_n_var = nf_nob i_n_var_;;
type i_num_var = (nf_nob,nf) i_num_var_;;

let hd_of_i_var =
 function
   `I ((v,_),_)
 | `Var (v,_) -> v

let hd_of =
 function
   `I ((v,_),_)
 | `Var(v,_) -> Some v
 | `N _ -> None
 | `Match _ -> assert false

let arity_of_hd =
function
  `I ((_,a),_)
| `Var(_,a) -> a
| _ -> 0 (* FIXME? *)

let lift m (t : nf) =
 let aux_var l (n, ar) = (if n < l then n else n+m), ar in
 let rec aux_i_num_var l =
  function
    `I(v,args) -> `I(aux_var l v, Listx.map (aux_nob l) args)
   | `Var v -> `Var(aux_var l v)
   | `N _ as x -> x
   | `Match(t,v,lift,bs,args) ->
      `Match(aux_i_num_var l t, v, lift + m, bs, List.map (aux_nob l) args)
 and aux_nob l =
  function
     #i_num_var as x -> (aux_i_num_var l x :> nf_nob)
   | `Lam(b,nf) -> `Lam (b, aux (l+1) nf)
   | `Pacman -> `Pacman
 and aux l =
  function
     #nf_nob as x -> (aux_nob l x :> nf)
   | `Bottom -> `Bottom
 in
  (aux 0 t : nf)
;;

(* put t under n lambdas, lifting t accordingtly *)
let rec make_lams t =
 function
   0 -> t
 | n when n > 0 -> `Lam (false, lift 1 (make_lams t (n-1)))
 | _ -> assert false

let free_vars' =
 let rec aux n = function
    `N _ -> []
  | `Var(x,ar) -> if x < n then [] else [(x-n,ar)]
  | `I((x,ar),args) ->
      (if x < n then [] else [(x-n,ar)]) @
      List.concat (List.map (aux n) (Listx.to_list args :> nf list))
  | `Lam(_,t) -> aux (n+1) t
  | `Match(t,_,liftno,bs,args) ->
      aux n (t :> nf) @
      List.concat (List.map (fun (_,t) -> aux (n-liftno) t) !bs) @
      List.concat (List.map (aux n) (args :> nf list))
  | `Bottom | `Pacman -> []
  in aux 0
;;
let free_vars = (List.map fst) ++ free_vars';;

module ToScott =
struct

let rec scott_of_nf = function
 | `N n -> Scott.mk_n n
 | `Var(v,_) -> Pure.V v
 | `Match(t,_,liftno,bs,args) ->
 let bs = List.map (fun (n,t) -> n, scott_of_nf (lift liftno (t :> nf))) !bs in
     let t = scott_of_nf (t :> nf) in
     let m = Scott.mk_match t bs in
     List.fold_left (fun acc t -> Pure.A(acc,scott_of_nf t)) m (args :> nf list)
 | `I((v,_), args) -> Listx.fold_left (fun acc t -> Pure.A(acc,scott_of_nf t)) (Pure.V v) (args :> nf Listx.listx)
 | `Lam(_,t) -> Pure.L (scott_of_nf t)
 | `Bottom -> Pure.B
 | `Pacman -> let f x = Pure.A (x,x) in f (Pure.L (Pure.L (f (Pure.V 0))))
end


(************ Pretty-printing ************************************)

(* let rec string_of_term l  = fun _ -> "";; *)

let rec string_of_term =
 let boundvar x = "v" ^ string_of_int x in
 let varname lev l n =
  if n < lev then boundvar (lev-n-1)
   else if n < List.length l then List.nth l (n-lev)
    else "`" ^ string_of_int (n-lev) in
 let rec string_of_term_w_pars lev l  = function
  | `Var(n,ar) -> varname lev l n ^ (if debug_display_arities then ":" ^ string_of_int ar else "")
  | `N n -> string_of_int n
  | `I _ as t -> "(" ^ string_of_term_no_pars_app lev l t ^ ")"
  | `Lam(_,`Bottom) -> "BOMB"
  | `Lam _ as t -> "(" ^ string_of_term_no_pars_lam lev l t ^ ")"
  | `Match(t,(v,ar),bs_lift,bs,args) ->
      (* assert (bs_lift = lev); *)
     "(["^ varname lev l v ^ (if debug_display_arities then ":"^ string_of_int ar else "") ^",match " ^ string_of_term_no_pars lev l (t :> nf) ^
     " with " ^ String.concat " | " (List.map (fun (n,t) -> string_of_int n ^ " => " ^ string_of_term l (t :> nf)) !bs) ^ "] " ^
     String.concat " " (List.map (string_of_term l) (args :> nf list)) ^ ")"
  | `Bottom -> "BOT"
  | `Pacman -> "PAC"
 and string_of_term_no_pars_app lev l  = function
  | `I((n,ar), args) -> varname lev l n ^ (if debug_display_arities then ":" ^ string_of_int ar else "") ^ " " ^ String.concat " " (List.map (string_of_term_w_pars lev l) (Listx.to_list args :> nf list))
  | #nf as t -> string_of_term_w_pars lev l t
 and string_of_term_no_pars_lam lev l  = function
  | `Lam(_,`Bottom) -> "BOMB"
  | `Lam(_,t) -> "λ" ^ boundvar lev ^ ". " ^ (string_of_term_no_pars_lam (lev+1) l t)
  | _ as t -> string_of_term_no_pars lev l t
 and string_of_term_no_pars lev l = function
  | `Lam _ as t -> string_of_term_no_pars_lam lev l t
  | #nf as t -> string_of_term_no_pars_app lev l t
 in string_of_term_no_pars 0
;;

let print ?(l=[]) = string_of_term l;;
let string_of_nf t = string_of_term [] (t :> nf);;

(************ Hereditary substitutions ************************************)

let cast_to_i_var =
 function
    #i_var as y -> (y : i_var)
  | t ->
    prerr_endline (print (t :> nf));
    assert false (* algorithm failed *)

let cast_to_i_n_var =
 function
    #i_n_var as y -> (y : i_n_var)
  | t ->
    prerr_endline (print (t :> nf));
    assert false (* algorithm failed *)

let cast_to_i_num_var =
 function
    #i_num_var as y -> (y : i_num_var)
  | t ->
    prerr_endline (print (t :> nf));
    assert false (* algorithm failed *)

let rec set_arity arity = function
(* FIXME because onlt variables should be in branches of matches, one day *)
| `Var(n,_) -> `Var(n,arity)
| `N _ | `Bottom | `Pacman as t -> t
| `Lam(false, t) -> `Lam(false, set_arity arity t)
| `Match(t,(n,_),bs_lift,bs,args) -> `Match(t,(n,arity),bs_lift,bs,args)
| `I _ | `Lam _ -> assert false

let minus1 n = if n = min_int then n else n - 1;;

let rec mk_app (h : nf) (arg : nf) =
 match arg with
  | `Bottom -> `Bottom
  | #nf_nob as arg ->
    match h with
     | `I(v,args) -> `I(v,Listx.append (Listx.Nil arg) args)
     | `Var v -> `I(v, Listx.Nil arg)
     | `Lam(truelam,nf) -> subst truelam true 0 arg (nf : nf) (* AC FIXME sanity check on arity *)
     | `Match(t,v,lift,bs,args) -> `Match(t,v,lift,bs,List.append args [arg])
     | `Bottom | `Pacman as t -> t
     | `N _ -> assert false (* Numbers cannot be applied *)
(*in let l = ["v0";"v1";"v2"] in
prerr_endline ("mk_app h:" ^ print ~l h ^ " arg:" ^ print ~l:l arg ^ " res:" ^ print ~l:l res); res*)

and mk_appl h args =
 (*prerr_endline ("MK_APPL: " ^ print h ^ " " ^ String.concat " " (List.map print args));*)
 List.fold_left mk_app h args

and mk_appx h args = Listx.fold_left mk_app h args

and mk_match t (n,ar) bs_lift bs args =
 (*prerr_endline ("MK_MATCH: ([" ^ print t ^ "] " ^ String.concat " " (Listx.to_list (Listx.map (fun (n,t) -> string_of_int n ^ " => " ^ print t) bs)) ^ ") " ^ String.concat " " (List.map print args));*)
 let m =
  match t with
    `N m as t ->
      (try
        let h = List.assoc m !bs in
        let h = set_arity (minus1 ar) h in
        let h = lift bs_lift h in
         h
       with Not_found ->
        `Match (t,(n,ar),bs_lift,bs,[]))
   (* We are assuming that the econding of matches is s.t.:
      - match PAC.. --> PAC
      - match BOT.. --> BOT *)
   | `Bottom -> `Bottom
   | `Pacman -> `Pacman
   | `Lam _ -> assert false
   | `I _ | `Var _ | `Match _ as t -> `Match(t,(n,ar),bs_lift,bs,[]) in
 mk_appl m args

and subst truelam delift_by_one what (with_what : nf(*_nob*)) (where : nf) =
 let rec aux_propagate_arity ar = function
 | `Lam(false, t) when not delift_by_one -> `Lam(false, aux_propagate_arity ar t)
 | `Match(`I(v,args),(x,_),liftno,bs,args') when not delift_by_one ->
    `Match(`I(v,args),(x,ar),liftno,bs,args')
 | `Var(i,oldar) -> `Var(i, if truelam then (assert (oldar = min_int); ar) else oldar)
 | _ as t -> t in
 let rec aux_i_num_var l =
  function
     `I((n,ar),args) ->
       if n = what + l then
       let args = Listx.map (aux l) (args :> nf Listx.listx) in
        mk_appx (lift l (aux_propagate_arity ar (with_what :> nf))) args
       else
        mk_appl (`Var ((if delift_by_one && n >= l then n-1 else n), ar)) (List.map (aux l) (Listx.to_list (args :> nf Listx.listx)))
   | `Var(n,ar) ->
       if n = what + l then
        lift l (aux_propagate_arity ar (with_what :> nf))
       else
        `Var((if delift_by_one && n >= l then n-1 else n), ar)
   | `N _ as x -> x
   | `Match(t,v,bs_lift,bs,args) ->
       let bs_lift = bs_lift + if delift_by_one then -1 else 0 in
       (* Warning! It now applies again the substitution in branches of matches.
          But careful, it does it many times, for every occurrence of
          the match. This is okay because what does not occur in with_what. *)
       let l' = l - bs_lift  in
       let with_what' = lift l' (with_what :> nf) in
       (* The following line should be the identity when delift_by_one = true because we
          are assuming the ts to not contain lambda-bound variables. *)
       bs := List.map (fun (n,t) -> n,subst truelam false what with_what' t) !bs ;
       let body = aux_i_num_var l t in
       mk_match body v bs_lift bs (List.map (aux l) (args :> nf list))
 and aux l(*lift*) =
(*function iii -> let res = match iii with*)
  function
   | #i_num_var as x -> aux_i_num_var l x
   | `Lam(b, nf) -> `Lam(b, aux (l+1) nf)
   | `Bottom -> `Bottom
   | `Pacman -> `Pacman
(*in let ll = ["v0";"v1";"v2"] in
prerr_endline ("subst l:" ^ string_of_int l ^ " delift_by_one:" ^ string_of_bool delift_by_one ^ " what:" ^ (List.nth ll what) ^ " with_what:" ^ print ~l:ll with_what ^ " where:" ^ print ~l:ll iii ^ " res:" ^ print ~l:ll res); res*)
 in
  aux 0 where
;;

(************ Parsing ************************************)

(************** Algorithm(s) ************************)

let eta_compare x y =
 (* let clex a b = let diff = ? a b in if diff = 0 then cont () else 0 in *)
 let clex aux1 aux2 (a1,a2) (b1,b2) =
  let diff = aux1 a1 b1 in if diff = 0 then aux2 a2 b2 else diff in
 let rec lex aux l1 l2 =
  match l1,l2 with
 | [], [] -> 0
 | [], _ -> -1
 | _, [] -> 1
 | x::xs, y::ys -> clex aux (lex aux) (x,xs) (y,ys) in
 let rec aux t1 t2 = match t1, t2 with
  | `Var(n,_) , `Var(m,_) -> compare n m
  | `I((n1,_), l1), `I((n2,_), l2) ->
    clex compare (lex aux) (n1, (Listx.to_list l1 :> nf list)) (n2, (Listx.to_list l2 :> nf list))
  | `Bottom, `Bottom
  | `Pacman, `Pacman -> 0
  | `Lam _, `N _ -> -1
  | `N _, `Lam _ -> 1
  | `Bottom, `Lam(_,t) -> -1
  | `Lam(_,t), `Bottom ->  1
  | `Lam(_,t1), `Lam(_,t2) -> aux t1 t2
  | `Lam(_,t1), t2 -> - aux t1 (mk_app (lift 1 t2) (`Var(0,-666)))
  | t2, `Lam(_,t1) ->   aux t1 (mk_app (lift 1 t2) (`Var(0,-666)))
  | `N n1, `N n2 -> compare n1 n2
  | `Match(u,_,bs_lift,bs,args), `Match(u',_,bs_lift',bs',args') ->
    let bs = List.sort (fun (n,_) (m,_) -> compare n m) !bs in
    let bs' = List.sort (fun (n,_) (m,_) -> compare n m) !bs' in
    clex aux (clex (lex (clex compare aux)) (lex aux)) ((u :> nf), (bs, (args :> nf list))) ((u' :> nf), (bs', (args' :> nf list)))
  | `Match _, _ -> -1
  | _, `Match _ -> 1
  | `N _, _ -> -1
  | _, `N _ -> 1
  | `I _, _ -> -1
  | _, `I _ -> 1
  | `Bottom, _ -> -1
  | _, `Bottom -> 1
  | `Pacman, _ -> -1
  | _, `Pacman -> 1
  in aux x y
;;

let eta_eq (#nf as x) (#nf as y) = 0 = eta_compare x y ;;

let rec eta_subterm sub t =
 if eta_eq sub t then true else
  match t with
  | `Lam(_,t') -> eta_subterm (lift 1 sub) t'
  | `Bottom
  | `Pacman  -> false
  | `Match(u,ar,liftno,bs,args) ->
     eta_subterm sub (u :> nf)
     || List.exists (fun (_, t) -> eta_subterm sub (lift liftno t)) !bs
     || List.exists (eta_subterm sub) (args :> nf list)
  | `I((v,_), args) -> List.exists (eta_subterm sub) ((Listx.to_list args) :> nf list) || (match sub with
   | `Var(v',_) -> v = v'
   | `I((v',_), args') -> v = v'
    && Listx.length args' < Listx.length args
    && List.for_all (fun (x,y) -> eta_eq x y) (List.combine (Util.take (Listx.length args') (Listx.to_list args :> nf list)) (Listx.to_list args' :> nf list))
   | _ -> false
   )
  | `N _ | `Var _ -> false
;;

let eta_subterm (#nf as x) (#nf as y) = eta_subterm x y;;

let max_arity_tms n =
 let max a b = match a, b with
  | None, None -> None
  | None, Some x
  | Some x, None -> Some x
  | Some x, Some y -> Some (Pervasives.max x y) in
 let aux_var l (m,a) = if n + l = m then Some a else None in
 let rec aux l = function
  | `Var v -> aux_var l v
  | `I(v,tms) -> max (aux_var l v) (aux_tms l (Listx.to_list tms :> nf list))
  | `Lam(_,t) -> aux (l+1) t
  | `Match(u,_,_,bs,args) -> max (max (aux l (u :> nf)) (aux_tms l (args :> nf list))) (aux_tms l (List.map snd !bs))
  | `N _ | `Bottom | `Pacman -> None
 and aux_tms l =
  List.fold_left (fun acc t -> max acc (aux l t)) None in
 fun tms -> aux_tms 0 (tms :> nf list)
;;