type t =
  | Var of string
  | Lam of string * t
  | App of t * t
  | Subst of string * t * t

module SSet = Set.Make (String)

let rec free_variables = function
  | Var x -> SSet.singleton x
  | Lam (x, b) -> SSet.remove x (free_variables b)
  | App (f, a) -> SSet.union (free_variables f) (free_variables a)
  | Subst (x, s, b) ->
      let fv_b = free_variables b in
      let fv_b' = SSet.remove x fv_b in
      if SSet.mem x fv_b then SSet.union fv_b' (free_variables s) else fv_b'

let free_variable_list t = SSet.elements (free_variables t)

let rec pretty = function
  | Var x -> x
  | Lam (x, b) -> "\\" ^ x ^ ". " ^ pretty b
  | App (f, a) ->
      let pf =
        match f with Var _ | App _ -> pretty f | _ -> "(" ^ pretty f ^ ")"
      in
      let pa = match a with Var _ -> pretty a | _ -> "(" ^ pretty a ^ ")" in
      pf ^ " " ^ pa
  | Subst (x, s, b) -> "[" ^ x ^ " := " ^ pretty s ^ "] " ^ pretty b

let fresh_name 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

let rename old nw t =
  let rec ren = function
    | Var z -> if z = old then Var nw else Var z
    | App (f, a) -> App (ren f, ren a)
    | Lam (z, b) -> if z = old then Lam (z, b) else Lam (z, ren b)
    | Subst (z, u, b) ->
        if z = old then Subst (z, ren u, b) else Subst (z, ren u, ren b)
  in
  ren t

let substitute x s body =
  let fv_s = free_variables s in
  let fv_b = free_variables body in
  let avoid0 = SSet.union fv_s fv_b in
  let rec sub avoid = function
    | Var y -> if y = x then s else Var y
    | App (f, a) -> App (sub avoid f, sub avoid a)
    | Lam (y, b) ->
        if y = x then Lam (y, b)
        else if (not (SSet.mem y fv_s)) || not (SSet.mem x (free_variables b)) then
          Lam (y, sub avoid b)
        else
          let y' = fresh_name avoid y in
          Lam (y', sub (SSet.add y' avoid) (rename y y' b))
    | Subst (y, u, b) ->
        if y = x then Subst (y, sub avoid u, b)
        else if (not (SSet.mem y fv_s)) || not (SSet.mem x (free_variables b)) then
          Subst (y, sub avoid u, sub avoid b)
        else
          let y' = fresh_name avoid y in
          Subst (y', sub avoid u, sub (SSet.add y' avoid) (rename y y' b))
  in
  sub avoid0 body

let step_limit = 10000

let rec elim_subst = function
  | Var x -> Var x
  | Lam (x, b) -> Lam (x, elim_subst b)
  | App (f, a) -> App (elim_subst f, elim_subst a)
  | Subst (x, s, b) -> substitute x (elim_subst s) (elim_subst b)

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

let rec step = function
  | App (Lam (x, b), a) -> Some (substitute x 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 (x, b) -> begin
      match step b with Some b' -> Some (Lam (x, b')) | None -> None
    end
  | _ -> None

let normalise t =
  let t = elim_subst t in
  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 _ -> ()
    | Lam (_, b) -> walk (d + 1) b
    | App (f, a) -> walk (d + 1) f; walk (d + 1) a
    | Subst (_, s, b) -> walk (d + 1) s; walk (d + 1) b
  in
  let t = elim_subst t 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 convert_from_named t = t
let convert_to_named t = t

let alpha_equivalent t1 t2 =
  let rec go d e1 e2 a b =
    match (a, b) with
    | Var x, Var y -> begin
        match (List.assoc_opt x e1, List.assoc_opt y e2) with
        | Some i, Some j -> i = j
        | None, None -> x = y
        | _ -> false
      end
    | Lam (x, b1), Lam (y, b2) -> go (d + 1) ((x, d) :: e1) ((y, d) :: e2) b1 b2
    | App (f1, a1), App (f2, a2) -> go d e1 e2 f1 f2 && go d e1 e2 a1 a2
    | _ -> false
  in
  go 0 [] [] (normalise t1) (normalise t2)

let alpha_equivalent_counted t1 t2 =
  let count = ref 0 in
  let maxd = ref 0 in
  let rec go d e1 e2 a b =
    incr count;
    if d > !maxd then maxd := d;
    match (a, b) with
    | Var x, Var y -> begin
        match (List.assoc_opt x e1, List.assoc_opt y e2) with
        | Some i, Some j -> i = j
        | None, None -> x = y
        | _ -> false
      end
    | Lam (x, b1), Lam (y, b2) -> go (d + 1) ((x, d) :: e1) ((y, d) :: e2) b1 b2
    | App (f1, a1), App (f2, a2) -> go d e1 e2 f1 f2 && go d e1 e2 a1 a2
    | _ -> false
  in
  let r = go 0 [] [] (normalise t1) (normalise t2) in
  (r, !count, !maxd)