type t =
  | BVar of int
  | FVar of string
  | Lam of t
  | App of t * t

module SSet = Set.Make (String)

exception Not_locally_closed

let rec free_variables = function
  | BVar _ -> SSet.empty
  | FVar x -> SSet.singleton x
  | Lam b -> free_variables b
  | App (f, a) -> SSet.union (free_variables f) (free_variables a)

let free_variable_list t = SSet.elements (free_variables t)

let rec pretty = function
  | BVar i -> "#" ^ string_of_int i
  | FVar x -> x
  | Lam b -> "\\. " ^ pretty b
  | App (f, a) ->
      let pf =
        match f with
        | BVar _ | FVar _ | App _ -> pretty f
        | _ -> "(" ^ pretty f ^ ")"
      in
      let pa =
        match a with BVar _ | FVar _ -> pretty a | _ -> "(" ^ pretty a ^ ")"
      in
      pf ^ " " ^ pa

let rec open_at k x = function
  | BVar i -> if i = k then FVar x else BVar i
  | FVar y -> FVar y
  | Lam b -> Lam (open_at (k + 1) x b)
  | App (f, a) -> App (open_at k x f, open_at k x a)

let open_top t x = open_at 0 x t

let rec close_at k x = function
  | BVar i -> BVar i
  | FVar y -> if y = x then BVar k else FVar y
  | Lam b -> Lam (close_at (k + 1) x b)
  | App (f, a) -> App (close_at k x f, close_at k x a)

let close_top t x = close_at 0 x t

let rec lc_at depth = function
  | BVar i -> i < depth
  | FVar _ -> true
  | Lam b -> lc_at (depth + 1) b
  | App (f, a) -> lc_at depth f && lc_at depth a

let lc t = lc_at 0 t

let require_lc t =
  if not (lc t) then raise Not_locally_closed

let rec subst_bvar j s = function
  | BVar i -> if i = j then s else if i > j then BVar (i - 1) else BVar i
  | FVar x -> FVar x
  | Lam b -> Lam (subst_bvar (j + 1) (shift s) b)
  | App (f, a) -> App (subst_bvar j s f, subst_bvar j s a)

and shift t =
  let rec go cutoff = function
    | BVar i -> if i >= cutoff then BVar (i + 1) else BVar i
    | FVar x -> FVar x
    | Lam b -> Lam (go (cutoff + 1) b)
    | App (f, a) -> App (go cutoff f, go cutoff a)
  in
  go 0 t

let rec subst_fvar x s = function
  | BVar i -> BVar i
  | FVar y -> if y = x then s else FVar y
  | Lam b -> Lam (subst_fvar x s b)
  | App (f, a) -> App (subst_fvar x s f, subst_fvar x s a)

let convert_from_named t =
  let rec go env = function
    | Named.Var x -> begin
        match List.assoc_opt x env with
        | Some i -> BVar i
        | None -> FVar x
      end
    | Named.Lam (x, b) ->
        Lam (go ((x, 0) :: List.map (fun (v, i) -> (v, i + 1)) env) b)
    | Named.App (f, a) -> App (go env f, go env a)
    | Named.Subst (x, s, b) ->
        go env (Named.substitute x s b)
  in
  go [] t

let convert_to_named t =
  require_lc t;
  let fv = free_variables t in
  let fresh_avoid avoid base =
    if not (SSet.mem base avoid) then base
    else
      let rec go i =
        let cand = base ^ string_of_int i in
        if SSet.mem cand avoid then go (i + 1) else cand
      in
      go 0
  in
  let rec go env = function
    | BVar i -> Named.Var (List.nth env i)
    | FVar x -> Named.Var x
    | Lam b ->
        let avoid =
          List.fold_left (fun a x -> SSet.add x a) fv env
        in
        let x = fresh_avoid avoid "v" in
        Named.Lam (x, go (x :: env) b)
    | App (f, a) -> Named.App (go env f, go env a)
  in
  go [] t

let is_value = function Lam _ -> true | _ -> false

let rec step = function
  | App (Lam b, a) -> Some (subst_bvar 0 a b)
  | App (f, a) when not (is_value f) -> begin
      match step f with
      | Some f' -> Some (App (f', a))
      | None -> begin
          match step a with Some a' -> Some (App (f, a')) | None -> None
        end
    end
  | App (f, a) -> begin
      match step a with Some a' -> Some (App (f, a')) | None -> None
    end
  | Lam b -> begin
      match step b with Some b' -> Some (Lam b') | None -> None
    end
  | _ -> None

let step_limit = 10000

let normalise t =
  require_lc t;
  let rec go n t =
    if n >= step_limit then t
    else match step t with Some t' -> go (n + 1) t' | None -> t
  in
  go 0 t

let normalise_counted t =
  require_lc t;
  let count = ref 0 in
  let maxd = ref 0 in
  let rec walk d t =
    incr count;
    if d > !maxd then maxd := d;
    match t with
    | BVar _ | FVar _ -> ()
    | Lam b -> walk (d + 1) b
    | App (f, a) -> walk (d + 1) f; walk (d + 1) a
  in
  walk 0 t;
  let rec go n t =
    if n >= step_limit then t
    else
      match step t with
      | Some t' -> walk 0 t'; go (n + 1) t'
      | None -> t
  in
  (go 0 t, !count, !maxd)

let substitute x s body = subst_fvar x s body

let alpha_equivalent t1 t2 =
  require_lc t1;
  require_lc t2;
  normalise t1 = normalise t2

let alpha_equivalent_counted t1 t2 =
  require_lc t1;
  require_lc t2;
  let count = ref 0 in
  let maxd = ref 0 in
  let rec go d a b =
    incr count;
    if d > !maxd then maxd := d;
    match (a, b) with
    | BVar i, BVar j -> i = j
    | FVar x, FVar y -> x = y
    | Lam b1, Lam b2 -> go (d + 1) b1 b2
    | App (f1, a1), App (f2, a2) -> go (d + 1) f1 f2 && go (d + 1) a1 a2
    | _ -> false
  in
  let r = go 0 (normalise t1) (normalise t2) in
  (r, !count, !maxd)