import numpy as np
from functools import partial
from math import sqrt, exp

def node(t,u):
    return u + t*(t+1)/2   

class BinomialTree(object):   
    def __init__(self,T):
        self.values = np.zeros(1 + node(T,T))
        self.numberOfPeriods = T    
    def __getitem__(self, n):
        return self.values[node(n[0],n[1])]        
    def __setitem__(self,n,value):
        self.values[node(n[0],n[1])] = value

def backRecurse(b, tAgg, ce):
    T = b.numberOfPeriods
    v = BinomialTree(T)
    for u in range(T+1):
        v[T,u] = b[T,u]    
    for t in range(T-1, -1, -1):
        for u in range(t+1):
            v[t,u] = tAgg(ce(v[t+1,u+1], v[t+1,u]) ,b[t,u])
    return v

def constVolTree(c0, T, U, D):
    c = BinomialTree(T)
    c[0,0] = c0
    for t in range(1, T+1):
        c[t,0]=c[t-1,0]*D
        for u in range(1, t+1):
            c[t,u] = c[t-1,u-1]*U
    return c

def siCE(uV, dV, p, cRRA): # assume cRRA away from one
    power = 1 - cRRA
    return (p*(uV**power) + (1-p)*(dV**power))**(1/power)

alphaValues = [2.0, 20.0, 40.0]
beta = 0.03
gamma = 2.0
delta = 1.0 / 1.5
muHi = 0.05
muLo = 0.00
sigma = 0.03
nPeriods = 520
horizon = 10.0
h = horizon / nPeriods
pHi = (1 + (muHi / sigma) * sqrt(h)) / 2
pLo = (1 + (muLo / sigma) * sqrt(h)) / 2
upGrowth = exp(sigma * sqrt(h))
c = constVolTree(100, nPeriods, upGrowth, 1 / upGrowth)
discount = exp(-beta * h)

tAgg = partial(siCE, p=discount, cRRA=delta)
hiCE = partial(siCE, p=pHi, cRRA = gamma)
loCE = partial(siCE, p=pLo, cRRA = gamma)

def ce(alpha):
    return lambda x,y: siCE(hiCE(x, y), loCE(x, y), 0.5, a)

for a in alphaValues:   
    u = backRecurse(c, tAgg, ce(a))
    print("For alpha = {0} the time-zero utility is {1}".format(a, u[0,0]))