import Control.Monad

alphaValues = [2.0, 20.0, 40.0]
beta = 0.03
gamma = 2.0
delta = 1/1.5
muHi = 0.05    -- optimistic drift
muLo = 0.00    -- pessimistic drift
sigma = 0.03
nPeriods = 520 :: Int -- total number of periods
h = 10.0 / (fromIntegral nPeriods)
c0  = 100.0     -- initial consumption
pHi = (1 + (muHi/sigma)*(sqrt h)) / 2
pLo = (1 + (muLo/sigma)*(sqrt h)) / 2
upGrowth = exp (sigma*(sqrt h))
dnGrowth = 1 / upGrowth
discount = exp (-beta * h)
siCE p cRRA (dn,up) = 
    if power == 1 
      then exp (q*(log dn) + p*(log up))
      else (q*(dn**power)+p*(up**power))**(1/power)
    where power = 1 - cRRA; q = 1 - p        
agg = siCE discount delta -- time aggregator
hiCE = siCE pHi gamma     -- optimistic EU CE
loCE = siCE pLo gamma     -- pessimistic EU CE

-- Consumption values at a given time are represented as a list,
-- with the head corresponding to least consumption (bottom of the tree)

-- Function to move forward one step on the consumption tree:
fwd xs = ((head xs)*dnGrowth):(map (*upGrowth) xs) 

-- Function to move backward one stepn the consumption tree:
bwd (x:xs)  = map (/upGrowth) xs 

-- Forward recursion to get terminal consumption values:
cVals acc t = if t <= 0 then acc else cVals (fwd acc) (t-1)
csLast = cVals [c0] (nPeriods-1) -- consumption at beginning of final period

-- Function computing time-zero utility as a function of alpha.
-- The tail-recursion uses accumulator acc = (cs,vs),
-- where cs lists beginning-of-period consumption 
-- and vs lists end-of-period continuation utility.
utility alpha = timeZeroUtility where
    soCE (x,y) = siCE 0.5 alpha (hiCE (x,y), loCE (x,y))
    backStep (cs,vs) = ( bwd cs, map agg $ zip cs $ map soCE $ zip vs (tail vs))
    uVals acc time2go = if time2go == 1 then snd (backStep acc) else uVals (backStep acc) (time2go - 1)                                
    [timeZeroUtility] = uVals (csLast, fwd csLast) nPeriods

main = do 
    forM_ alphaValues (\a -> do
        let alpha = show a
        let util  = show (utility a)
        putStrLn $ "For alpha = " ++ alpha ++ " the time-zero utility is " ++ util)