object exp {


case class Scope(m: Map[Var[_], Exp[_]]) {
  def put[T](pair: (Var[T], Exp[T])): Scope = new Scope(m + pair)
  def apply[T](v: Var[T]): Either[String, Exp[T]] = m.get(v) match {
    case Some(exp) => Right(exp.asInstanceOf[Exp[T]])
    case None => Left(s"$v not found. Scope: $m")
  }
}
object Scope {
  def empty: Scope = Scope(Map.empty)
}

sealed trait Exp[A] {
  /**
   * Either evaluate the Expression, or return an error string
   * (about missing Vars is the only thing that can go wrong)
   */
  private[exp] def run(scope: Scope): Either[String, A]
}
case class Const[A](get: A) extends Exp[A] {
  private[exp] def run(scope: Scope) = Right(get)
}
case class Pair[A, B](first: Exp[A], second: Exp[B]) extends Exp[(A, B)] {
  private[exp] def run(scope: Scope) = for {
    a <- first.run(scope).right
    b <- second.run(scope).right
  } yield (a, b)
}
case class Apply[B, A](in: Exp[B], fn: Exp[Fn[B, A]]) extends Exp[A] {
  private[exp] def run(scope: Scope) = for {
    b <- in.run(scope).right
    f <- fn.run(scope).right
  } yield f.run(b)
}
case class Var[A](name: String) extends Exp[A] {
  private[exp] def run(scope: Scope) = for {
    a <- scope(this).right
    result <- a.run(scope).right
  } yield result
}
case class Lambda[A, B](x: Var[A], result: Exp[B]) extends Exp[Fn[A, B]] {
  private[exp] def run(scope: Scope) = Right(Fn.LambdaFn(this, scope))
}
// Private Fn class, so we can't write general scala functions
sealed trait Fn[A, B] {
  private[exp] def run(a: A): B
}

object Exp {
  // would not be here in real life, as it is an escape hatch
  // instead, we would convert the Exp to code/source that can
  // be run elsewhere
  def run[A](e: Exp[A]): Either[String, A] = e.run(Scope.empty)

  implicit class FnExp[A, B](val e: Exp[Fn[A, B]]) extends AnyVal {
    def apply(a: Exp[A]): Exp[B] = Apply[A, B](a, e)
    def andThen[C](that: Exp[Fn[B, C]]): Exp[Fn[A, C]] = {
      val va = Var[A]("a")
      val eb = Apply(va, e)
      val ec = Apply(eb, that)
      Lambda(va, ec)
    }
    def zip[C](that: Exp[Fn[A, C]]): Exp[Fn[A, (B, C)]] = {
      val va = Var[A]("a")
      val b = Apply(va, e)
      val c = Apply(va, that)
      Lambda(va, Pair(b, c))
    }
  }
}

object Fn {
  def fst[A, B](p: Exp[(A, B)]): Exp[A] =
    Apply(p, lift(Fst[A, B]()))
  def snd[A, B](p: Exp[(A, B)]): Exp[B] =
    Apply(p, lift(Snd[A, B]()))
  def addInt(a: Exp[Int], b: Exp[Int]): Exp[Int] =
    Apply(Pair(a, b), lift(AddInt))

  def ifThenElse[A](b: Exp[Boolean], iftrue: Exp[A], iffalse: Exp[A]): Exp[A] =
    Apply(Pair(b, Pair(iftrue, iffalse)), lift(IfThenElse[A]()))

  // non fundamental ops
  def lift[A, B](f: Fn[A, B]): Exp[Fn[A, B]] = Const(f)
  def swap[A, B](p: Exp[(A, B)]): Exp[(B, A)] =
    Pair(snd(p), fst(p))

  def curry[A, B, C](fn: Exp[Fn[(A, B), C]]): Exp[Fn[A, Fn[B, C]]] = {
    val a = Var[A]("A")
    val b = Var[B]("B")
    val pair = Pair(a, b)
    val c = Apply(pair, fn)
    val fnbc = Lambda(b, c)
    Lambda(a, fnbc)
  }
  def uncurry[A, B, C](fn: Exp[Fn[A, Fn[B, C]]]): Exp[Fn[(A, B), C]] = {
    val pair = Var[(A, B)]("pair")
    val a = fst(pair)
    val fnbc = Apply(a, fn)
    val b = snd(pair)
    val c = Apply(b, fnbc)
    Lambda(pair, c)
  }

  // implementations
  private[exp] case class LambdaFn[A, B](l: Lambda[A, B], scope: Scope) extends Fn[A, B] {
    def run(a: A) = {
      val newScope = scope.put(l.x -> Const(a))
      val b = l.result.run(newScope)
      b.right.get
    }
  }
  private case class Fst[A, B]() extends Fn[(A, B), A] {
    def run(b: (A, B)) = b._1
  }
  private case class Snd[A, B]() extends Fn[(A, B), B] {
    def run(b: (A, B)) = b._2
  }
  private case object AddInt extends Fn[(Int, Int), Int] {
    def run(p: (Int, Int)) = p._1 + p._2
  }
  private case class IfThenElse[A]() extends Fn[(Boolean, (A, A)), A] {
    def run(in: (Boolean, (A, A))) = if(in._1) in._2._1 else in._2._2
  }
}

}

/**
in the repl:
scala> :load expression.scala
Loading expression.scala...
defined module exp

scala> import exp._
import exp._

scala> Fn.ifThenElse(Const(true), Const(1), Const(2))
res16: exp.Exp[Int] = Apply(Pair(Const(true),Pair(Const(1),Const(2))),Const(IfThenElse()))

scala> Exp.run(res16)
res17: Either[String,Int] = Right(1)

scala> Fn.ifThenElse(Const(false), Const(1), Const(2))
res18: exp.Exp[Int] = Apply(Pair(Const(false),Pair(Const(1),Const(2))),Const(IfThenElse()))

scala> Exp.run(res18)
res19: Either[String,Int] = Right(2)

scala> Exp.run(Var("yo"))
res21: Either[String,Nothing] = Left(Var(yo) not found. Scope: Map())
*/