type t =
  | Var of int
  | Free of string
  | Lam of t
  | App of t * t

module SSet = Set.Make (String)

exception Unbound_index of int

let rec free_variables = function
  | Var _ -> SSet.empty
  | Free 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
  | Var i -> "#" ^ string_of_int i
  | Free x -> x
  | Lam b -> "\\. " ^ pretty b
  | App (f, a) ->
      let pf =
        match f with Var _ | Free _ | App _ -> pretty f | _ -> "(" ^ pretty f ^ ")"
      in
      let pa =
        match a with Var _ | Free _ -> pretty a | _ -> "(" ^ pretty a ^ ")"
      in
      pf ^ " " ^ pa

let rec shift cutoff amount t =
  match t with
  | Var i -> if i >= cutoff then Var (i + amount) else Var i
  | Free x -> Free x
  | Lam b -> Lam (shift (cutoff + 1) amount b)
  | App (f, a) -> App (shift cutoff amount f, shift cutoff amount a)

let rec subst j s = function
  | Var i ->
      if i = j then s else if i > j then Var (i - 1) else Var i
  | Free x -> Free x
  | Lam b -> Lam (subst (j + 1) (shift 0 1 s) b)
  | App (f, a) -> App (subst j s f, subst j s a)

let substitute j s body = subst j s body

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

let convert_from_named t = convert_from_named_depth [] t

let convert_to_named t =
  let fv = free_variables t in
  let fresh_avoid avoid base =
    if not (List.mem base avoid) then base
    else
      let rec go i =
        let cand = base ^ string_of_int i in
        if List.mem cand avoid then go (i + 1) else cand
      in
      go 0
  in
  let rec go env = function
    | Var i -> begin
        match List.nth_opt env i with
        | Some x -> Named.Var x
        | None -> raise (Unbound_index i)
      end
    | Free x -> Named.Var x
    | Lam b ->
        let avoid = SSet.elements 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 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 =
  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 =
  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
    | Var _ | Free _ -> ()
    | 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 alpha_equivalent t1 t2 = normalise t1 = normalise t2

let rec well_scoped depth = function
  | Var i -> i < depth
  | Free _ -> true
  | Lam b -> well_scoped (depth + 1) b
  | App (f, a) -> well_scoped depth f && well_scoped depth a

let alpha_equivalent_counted t1 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
    | Var i, Var j -> i = j
    | Free x, Free 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)