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 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 rec rename old nw = function
  | Var z -> if z = old then Var nw else Var z
  | App (f, a) -> App (rename old nw f, rename old nw a)
  | Lam (z, b) -> if z = old then Lam (z, b) else Lam (z, rename old nw b)
  | Subst (z, u, b) ->
      if z = old then Subst (z, rename old nw u, b)
      else Subst (z, rename old nw u, rename old nw b)

let step_limit = 10000

let push_subst x s body =
  let fv_s = free_variables s in
  let rec go = function
    | Var y -> if y = x then s else Var y
    | App (t, u) -> App (go t, go u)
    | Lam (y, b) ->
        if y = x then Lam (y, b)
        else if not (SSet.mem y fv_s) then Lam (y, go b)
        else
          let avoid =
            SSet.union fv_s
              (SSet.union (free_variables b) (free_variables s))
          in
          let y' = fresh avoid y in
          Lam (y', go (rename y y' b))
    | Subst (y, u, b) ->
        if y = x then Subst (y, go u, b)
        else if not (SSet.mem y fv_s) then Subst (y, go u, go b)
        else
          let avoid =
            SSet.union fv_s
              (SSet.union (free_variables b) (free_variables s))
          in
          let y' = fresh avoid y in
          Subst (y', go u, go (rename y y' b))
  in
  go body

let rec reduce = function
  | Var x -> Var x
  | Lam (x, b) -> Lam (x, reduce b)
  | App (f, a) -> begin
      let f' = reduce f in
      let a' = reduce a in
      match f' with Lam (x, b) -> reduce (push_subst x a' b) | _ -> App (f', a')
    end
  | Subst (x, s, b) -> reduce (push_subst x (reduce s) b)

let normalise t =
  let rec go n t =
    if n >= step_limit then t
    else
      let t' = reduce t in
      if t' = t then t else go (n + 1) t'
  in
  go 0 t

let rec to_named = function
  | Var x -> Named.Var x
  | Lam (x, b) -> Named.Lam (x, to_named b)
  | App (f, a) -> Named.App (to_named f, to_named a)
  | Subst (x, s, b) -> Named.Subst (x, to_named s, to_named b)

let rec of_named = function
  | Named.Var x -> Var x
  | Named.Lam (x, b) -> Lam (x, of_named b)
  | Named.App (f, a) -> App (of_named f, of_named a)
  | Named.Subst (x, s, b) -> Subst (x, of_named s, of_named b)

let alpha_equivalent t1 t2 =
  Named.alpha_equivalent (to_named (normalise t1)) (to_named (normalise t2))

let alpha_equivalent_susp 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
    | Subst (x, s1, b1), Subst (y, s2, b2) ->
        go d e1 e2 s1 s2
        && go (d + 1) ((x, d) :: e1) ((y, d) :: e2) b1 b2
    | Subst _, _ -> go d e1 e2 (normalise a) b
    | _, Subst _ -> go d e1 e2 a (normalise b)
    | _ -> false
  in
  go 0 [] [] t1 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
    | Subst (x, s1, b1), Subst (y, s2, b2) ->
        go d e1 e2 s1 s2
        && go (d + 1) ((x, d) :: e1) ((y, d) :: e2) b1 b2
    | Subst _, _ ->
        let a' = normalise a in
        if a' = a then false else go d e1 e2 a' b
    | _, Subst _ ->
        let b' = normalise b in
        if b' = b then false else go d e1 e2 a b'
    | _ -> false
  in
  let r = go 0 [] [] t1 t2 in
  (r, !count, !maxd)

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
  walk 0 t;
  let rec go n t =
    if n >= step_limit then t
    else
      let t' = reduce t in
      walk 0 t';
      if t' = t then t else go (n + 1) t'
  in
  (go 0 t, !count, !maxd)

let substitute x s body = Subst (x, s, body)

let convert_from_named t = of_named t
let convert_to_named t = to_named (normalise t)