// DEFINE A TYPE FOR HOLDING A RECOMBINING BINOMIAL TREE
type BinomialTree (T : int) =
    static let Node t u =  u + t*(t+1)/2 // t = time, u = # ups
    let N = 1 + Node T T
    let values : float array = Array.zeroCreate N
    member __.Item
        with get(t,u) = values.[Node t u]
        and  set(t,u) value = values.[Node t u] <- value
    member __.Periods = T
    member __.Nodes = N

// FUNCTION FOR CONSTRUCTING A CONSTANT VOLATILITY TREE
let ConstantVolTree c0 T U D =
    let C = BinomialTree(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:BinomialTree) tAgg CE =
    let T = C.Periods
    let U = BinomialTree (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 = 120   // 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 siCE p cRRA uV dV =
    match cRRA with
    | 1.0 -> exp(p*log(uV)+(1.0 - p)*log(dV))
    | _   -> 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]