Henk.scala

object Henk:
  def removeFirst[A](v: A, t: List[A]): List[A] = t match
    case Nil => Nil
    case x :: xs => if x == v
                    then xs
                    else x :: removeFirst(v, xs)

  enum Term:
    case Var(v: String)
    case Sort
    case Kind
    case Abs(v: String, t: Term, b: Term)
    case Pi(v: String, t: Term, b: Term)
    case App(a: Term, b: Term)

    type Env = List[(String, Term)]

    def typ(g: Env): Option[Term] = this match
      case Kind => Some(Sort)
      case Var(x) =>
        g.find { case (x, _) => true } match
          case Some((_, y)) => Some(y)
          case None => None
      case App(f, a) =>
        val tf = f.typ(g)
        val ta = a.typ(g)
        (tf, ta) match
          case (Some(Pi(x, tx, tb)), Some(ta)) =>
            if eqBeta(tx, ta)
            then Some(tb.substi(x, a))
            else None
          case _ => None
      case Abs(x, tx, b) =>
        b.typ((x, tx) :: g) match
          case Some(tb) =>
            val te = Pi(x, tx, tb)
            if te.typ(g).isDefined
            then Some(te)
            else None
          case None => None
      case Pi(x, tx, tb) =>
        (tx.redTyp(g), tb.redTyp((x, tx) :: g)) match
          case (Some(s), Some(t)) =>
            if to(s, t)
            then Some(t)
            else None
          case _ => None
      case _ => None

    def redTyp(g: Env): Option[Term] =
      this.typ(g) match
        case Some(v) => v.red()
        case _ => None

    def eqBeta(ta: Term, tb: Term): Boolean =
      (ta.red(), tb.red()) match
        case (Some(x), Some(y)) => eqAlpha(x, y, Nil, Nil)
        case _ => false

    def red(): Option[Term] = this match
      case App(Abs(x, tx, b), a) => b.substi(x, a).red()
      case _ => Some(this)

    def eqAlpha(a: Term, b: Term, sa: List[String], sb: List[String]): Boolean =
      (a, b) match
        case (Kind, Kind) => true
        case (Sort, Sort) => true
        case (Var(x), Var(y)) =>
          sa.indexOf(x) == sb.indexOf(y)
        case (App(a1, b1), App(a2, b2)) =>
          eqAlpha(a1, a2, sa, sb) && eqAlpha(b1, b2, sa, sb)
        case (Abs(x, a1, b1), Abs(y, a2, b2)) =>
          eqAlpha(a1, a2, x :: sa, y :: sb) && eqAlpha(b1, b2, x :: sa, y :: sb)
        case (Pi(x, a1, b1), Pi(y, a2, b2)) =>
          eqAlpha(a1, a2, x :: sa, y :: sb) && eqAlpha(b1, b2, x :: sa, y :: sb)
        case (_, _) => false

    def fv(): List[String] = this match
      case Kind => Nil
      case Sort => Nil
      case Var(x) => List(x)
      case App(f, a) => f.fv() ++ a.fv()
      case Abs(y, ty, b) => ty.fv() ++ removeFirst(y, b.fv())
      case Pi(y, ty, b) => ty.fv() ++ removeFirst(y, b.fv())

    def substi(x: String, s: Term): Term = this match
      case Kind => Kind
      case Sort => Sort
      case Var(y) => if (y == x) then s else this
      case App(f, a) =>
        App(f.substi(x, s), a.substi(x, s))
      case Abs(y, ty, b) if x == y =>
        Abs(y, ty.substi(x, s), b)
      case Abs(y, ty, b) if s.fv().contains(y) =>
        val z = y ++ "_"
        Abs(z, ty.substi(x, s), b.substiVar(y, z).substi(x, s))
      case Abs(y, ty, b) =>
        Abs(y, ty.substi(x, s), b.substi(x, s))
      case Pi(y, ty, b) if x == y =>
        Pi(y, ty.substi(x, s), b)
      case Pi(y, ty, b) if s.fv().contains(y) =>
        val z = y ++ "_"
        Pi(z, ty.substi(x, s), b.substiVar(y, z).substi(x, s))
      case Pi(y, ty, b) =>
        Pi(y, ty.substi(x, s), b.substi(x, s))

    def substiVar(x: String, z: String): Term = this match
      case Kind => Kind
      case Sort => Sort
      case Var(y) if y == x => Var(z)
      case Var(_) => this
      case App(f, a) =>
        App(f.substiVar(x, z), a.substiVar(x, z))
      case Abs(y, ty, b) =>
        Abs(y, ty.substiVar(x, z), b.substiVar(x, z))
      case Pi(y, ty, b) =>
        Pi(y, ty.substiVar(x, z), b.substiVar(x, z))

    def to(a: Term, b: Term): Boolean = (a, b) match
      case (Kind, Kind) => true
      case (Kind, Sort) => true
      case (Sort, Kind) => true
      case (Sort, Sort) => true
      case (_, _) => false

  val test: Term =
    Term.App(
      Term.Abs("x", Term.Pi("x", Term.Kind, Term.Kind), Term.Var("x")),
      Term.Abs("y", Term.Kind, Term.Var("y"))
    )

  def main(args: Array[String]) =
    println(test.toString)
    println("type:")
    println(test.typ(Nil).toString)
    println("result:")
    println(test.red().toString)