object RecursiveUtilityExample extends App {

  import scala.math._

  val alphaValues = List(2, 20, 40)
  val beta = 0.03
  val gamma = 2.0
  val delta = 1.0 / 1.5
  val muHi = 0.05
  val muLo = 0.00
  val sigma = 0.03
  val nPeriods = 120
  val horizon = 10.0
  val h = horizon / nPeriods
  val pHi = (1 + (muHi / sigma) * sqrt(h)) / 2
  val pLo = (1 + (muLo / sigma) * sqrt(h)) / 2
  val upGrowth = exp(sigma * sqrt(h))
  val C = BinomialTree.constVol(100, nPeriods, upGrowth, 1 / upGrowth)
  val discount = exp(-beta * h)
  val tAgg = siCE(discount, delta, _: Double, _: Double)
  val hiCE = siCE(pHi, gamma, _: Double, _: Double)
  val loCE = siCE(pLo, gamma, _: Double, _: Double)
  for (alpha <- alphaValues) {
    val U = C.backRecurse(tAgg, (x, y) => siCE(0.5, alpha, hiCE(x, y), loCE(x, y)))
    println("For alpha = " + alpha + " the time-zero utility is " + U(0, 0))
  }

  def siCE(p: Double, cRRA: Double, uV: Double, dV: Double) = {
    require(p > 0 && p < 1)
    val q = 1 - p
    if (cRRA == 1)
      exp(p * log(uV) + q * log(dV))
    else {
      val power = 1 - cRRA
      pow(p * pow(uV, power) + q * pow(dV, power), 1 / power)
    }
  }
}

class BinomialTree private (T: Int) {

  import BinomialTree.node

  private val values = new Array[Double](1 + node(T, T))

  def apply(t: Int, u: Int) = {
    require(t >= u && u >= 0)
    values(node(t, u))
  }
  def update(t: Int, u: Int, value: Double) {
    require(t >= u && u >= 0)
    values(node(t, u)) = value
  }
  def backRecurse(tAgg: (Double, Double) => Double, CE: (Double, Double) => Double) = {
    val V = BinomialTree(T)
    for (u <- 0 to T) V(T, u) = this(T, u)
    for (t <- T - 1 to 0 by -1)
      for (u <- 0 to t)
        V(t, u) = tAgg(CE(V(t + 1, u + 1), V(t + 1, u)), this(t, u))
    V
  }
}

object BinomialTree {

  def apply(T: Int) = { require(T > 0); new BinomialTree(T) }

  private def node(t: Int, u: Int): Int = u + t * (t + 1) / 2

  def constVol(c0: Double, T: Int, U: Double, D: Double) = {
    val C = BinomialTree(T)
    C(0, 0) = c0
    for (t <- 1 to T) {
      C(t, 0) = C(t - 1, 0) * D
      for (u <- 1 to t) C(t, u) = C(t - 1, u - 1) * U
    }
    C
  }
}