package binomial

object functionalRecursions extends App {

  import scala.math._

  val alphaValues = List(2.0, 20.0, 40.0)
  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 dnGrowth = 1 / upGrowth
  val discount = exp(-beta * h)
  val c0 = 100.0

  // function defining scale invariant expected utility CE
  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)
    }
  }
  val tAgg = siCE(discount, delta) _ tupled // time aggregator
  val hiCE = siCE(pHi, gamma) _ tupled // optimistic expected utility CE
  val loCE = siCE(pLo, gamma) _ tupled // pessimistic expected utility CE

  // Consumption at a given time is represented as a list
  // whose head is at the bottom of the tree (least consumption).
  // The following function takes such a list and gives the next one.
  def fwd(xs: List[Double]): List[Double] = (xs.head * dnGrowth) :: (xs map (_ * upGrowth))
  // The entire consumption tree is a list of lists, whose head is terminal consumption:
  val constVolTree: List[List[Double]] =
    (1 to nPeriods).toList.foldLeft(List(List(c0)))((css, _) => fwd(css.head) :: css)
 
  // The following function reduces the tree by performing one step of the back recursion:
  def bwd(alpha: Double, css: List[List[Double]]): List[List[Double]] =
    css match {
      case _ :: Nil => css
      case vs :: cs :: rest => {
        val ces = vs.tail zip vs map (u => siCE(0.5, alpha)(hiCE(u), loCE(u)))
        val uts = ces zip cs map tAgg
        uts :: rest
      }
    }
  // compute time-zero utility as a function of the parameter alpha
  def utility(alpha: Double) = {
       val List(List(timeZeroUtility)) =
         (1 to nPeriods).toList.foldLeft(constVolTree)((css, _) => bwd(alpha, css))
       timeZeroUtility
  }
  
  // report results
  for (alpha <- alphaValues) {
    println("For alpha = " + alpha + " the time-zero utility is " + utility(alpha))
  }

}