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 i_var_ = [ `I of var * 'nf Listx.listx | `Var of var ]
type 'nf i_n_var_ = [ `N of int | 'nf i_var_ ]
type 'nf i_num_var_ = [
  | 'nf i_n_var_
  | `Match of 'nf i_num_var_ * (* originating var *) var * (*lift*) int * (*branches*)(int * 'nf) list ref * (*args*)'nf list
]
type 'nf nf_ = [ `Lam of (* was_unpacked *) bool * 'nf nf_ | 'nf i_num_var_ ]
type nf = nf nf_
type i_var = nf i_var_;;
type i_n_var = nf i_n_var_;;
type i_num_var = 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 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 l) args)
 and aux l =
  function
     #i_num_var as x -> (aux_i_num_var l x :> nf)
   | `Lam(b,nf) -> `Lam (b, aux (l+1) nf)
 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))
  | `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)
  in aux 0
;;
let free_vars = (List.map fst) ++ free_vars';;

module ToScott =
struct

let delta = let open Pure in L(A(V 0, V 0))

let bomb = ref(`Var(-1, -666));;

let rec t_of_i_num_var =
 function
  | `N n -> Scott.mk_n n
  | `Var(v,_) as x -> assert (x <> !bomb); Pure.V v
  | `Match(t,_,liftno,bs,args) ->
      let bs = List.map (
        function (n,t) -> n,
         (if t = !bomb then delta
          else Pure.L (t_of_nf (lift (liftno+1) t)))
        ) !bs in
      let t = t_of_i_num_var t in
      let m = Scott.mk_match t bs in
      let m = Pure.A(m,delta) in
      List.fold_left (fun acc t -> Pure.A(acc,t_of_nf t)) m args
  | `I((v,_), args) -> Listx.fold_left (fun acc t -> Pure.A(acc,t_of_nf t)) (Pure.V v) args
and t_of_nf =
 function
  | #i_num_var as x -> t_of_i_num_var x
  | `Lam(b,f) -> Pure.L (t_of_nf f)

end


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

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

let 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 - lev < 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 _ as t -> "(" ^ string_of_term_no_pars_lam lev l t ^ ")"
  | `Match(t,(v,ar),bs_lift,bs,args) ->
      (* assert (bs_lift = lev); *)
     "(["^ varname 0 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)) ^ ")"
 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(_,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
 and string_of_term t = string_of_term_no_pars 0 t in
 string_of_term
;;

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 _ 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) =
(*let res =*)
 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])
  | `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));*)
 match t with
    `N m ->
      (try
        let h = List.assoc m !bs in
        let h = set_arity (minus1 ar) h in
        let h = lift bs_lift h in
        mk_appl h args
       with Not_found ->
        `Match (t,(n,ar),bs_lift,bs,args))
  | `I _ | `Var _ | `Match _ -> `Match(t,(n,ar),bs_lift,bs,args)

and subst truelam delift_by_one what (with_what : nf) (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
        mk_appx (lift l (aux_propagate_arity ar with_what)) (Listx.map (aux l) args)
       else
        `I (((if delift_by_one && n >= l then n-1 else n), ar), Listx.map (aux l) args)
   | `Var(n,ar) ->
       if n = what + l then
        lift l (aux_propagate_arity ar with_what)
       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 = cast_to_i_num_var (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)
(*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
;;

(************** 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) (n2, Listx.to_list l2)
  | `Lam _, `N _ -> -1
  | `N _, `Lam _ -> 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)) ((u' :> nf), (bs', args'))
  | `Match _, _ -> -1
  | _, `Match _ -> 1
  | `N _, _ -> -1
  | _, `N _ -> 1
  | `I _, _ -> -1
  | _, `I _ -> 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'
  | `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)) (Listx.to_list args'))
   | _ -> 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))
  | `Lam(_,t) -> aux (l+1) t
  | `Match(u,_,_,bs,args) -> max (max (aux l (u :> nf)) (aux_tms l args)) (aux_tms l (List.map snd !bs))
  | `N _ -> 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)
;;

let get_first_args var =
let rec aux l = function
| `Lam(_,t) -> aux (l+1) t
| `Match(u,orig,liftno,bs,args) -> Util.concat_map (aux l) args
| `I((n,_), args) -> if n = var + l then [Listx.last args] else []
| `N _
| `Var _ -> []
in aux 0
;;

let compute_arities m =
 let rec aux n tms =
  if n = 0
   then []
   else
    let tms = Util.filter_map (function `Lam(_,t) -> Some t | _ -> None ) tms in
    let arity = match max_arity_tms (m-n) tms with None -> -666 | Some x -> x in
     arity :: (aux (n-1) tms)
 in fun tms -> List.rev (aux m tms)
;;

let compute_arities var special_k all_tms =
 let tms = List.fold_left (fun acc t -> acc @ (get_first_args var t)) [] all_tms in
 compute_arities special_k tms
;;