module Binomial

// Function to make a binomial tree with T periods (times 0,1,...,T) as an array of arrays 
let make_tree T = Array.init (T+1) (fun i -> Array.create (i+1) 0.0)

// Function for constructing a constant volatility tree
let ConstantVolTree c0 T U D =
    let C = make_tree T
    C.[0].[0] <- c0
    for t = 1 to T do
        C.[t].[0] <- C.[t-1].[0]*D
        for u = 1 to t do
            C.[t].[u] <- C.[t-1].[u-1]*U
    C

// Function for computing a general backward recursion
let BackwardRecursiveTree (C: float [][]) tAgg CE =
    let T = (Array.length C) - 1
    let U = make_tree T
    for u = 0 to T do U.[T].[u] <- C.[T].[u]
    for t = T-1 downto 0 do
        for u = 0 to t do
            U.[t].[u] <- tAgg (CE U.[t+1].[u+1] U.[t+1].[u]) C.[t].[u]
    U

// DEFINE THE PARAMETERS
let alphaValues = [2.0;20.0;40.0] //list of values of alpha
let beta   =  0.03   // rate of impatience
let gamma  =  2.0    // first-order CRRA
let delta  =  1.0/1.5     // inverse EIS
let muHi   =  0.05   // high consumption growth scenario
let muLo   =  0.00   // low consumption growth scenario
let sigma  =  0.03   // known log-consumption volatility
let nPeriods = 3600   // total number of periods = 10 / h
let h = 10.0 / (float nPeriods) // time length per period = T / Nperiods
let pHi = (1.0 + (muHi/sigma)*(sqrt h)) / 2.0  // optimistic up probability
let pLo = (1.0 + (muLo/sigma)*(sqrt h)) / 2.0  // pessimistic up probability
let upGrowth = exp (sigma*(sqrt h))  // up consumption growth
let dnGrowth = 1.0 / upGrowth  // down consumption growth
let discount = exp(-beta * h)  //discount rate over one period

// CONSTRUCT THE CONSUMPTION TREE
let C = ConstantVolTree 100.0 nPeriods upGrowth dnGrowth

// DEFINE SCALE INVARIANT CE AND TIME AGGREGATOR   
let inline siCE p cRRA uV dV = // assume cRRA away from one
    let power = 1.0 - cRRA
    (p*(uV**power)+(1.0 - p)*(dV**power))**(1.0/power)    
let tAgg = siCE discount delta //current consumption corresponds to down state
let hiCE = siCE pHi gamma
let loCE = siCE pLo gamma

// COMPUTE TIME-ZERO UTILITY FOR VARIOUS VALUES OF ALPHA
for alpha in alphaValues do
    let CE x y = siCE 0.5 alpha (hiCE x y) (loCE x y)
    let U = BackwardRecursiveTree C tAgg CE
    printfn "For alpha = %f the time-zero utility is %f" alpha U.[0].[0]