import Control.Monad
import Data.List

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 -- assumes cRRA = 1 or is away from 1
  | cRRA == 1   = exp (q * log dn + p * log up)
  | otherwise   = (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
soCE alpha dn up = siCE 0.5 alpha (hiCE dn up) (loCE dn up) -- second-order EU CE

-- Time-t consumption is a list ordered from lowest to highest consumption scenarios.
-- The consumption tree is the list of lists:
-- [time-T consumption, time-(T-1) consumption, ... , time-0 consumption]

-- Function to move forward in time one step on the consumption tree:
fwd :: [Double] -> [Double]
fwd xs = head xs * dnGrowth : map (*upGrowth) xs  

-- Construct the constant volatility consumption tree
constVolTree :: [[Double]]
constVolTree = foldl' (\css _ -> fwd (head css) : css) [[c0]] [1..nPeriods]

-- Computation representing the back-recursion on the consumption tree as a function of alpha
-- The terminal consumption (head of list) is assumed equal to the continuation utility
bwd :: Double -> [[Double]] -> [[Double]]
bwd alpha css =
  case css of
      _:[] -> css
      vs:cs:rest -> uts:rest where
        ces = zipWith (soCE alpha) vs (tail vs)
        uts = zipWith agg cs ces

timeZeroUtility alpha = u where
  [[u]] = foldl' (\css _ -> bwd alpha css) constVolTree [1..nPeriods]


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

-- to compile, type: ghc --make -O2 funcRec.hs