```module TreeGrowth where

import qualified Probability
import Probability
(Dist, R, Space, mapD, normal, unfoldT, certainly, printR,
Trans, RTrans, Expand, RExpand, (*.), (*..), (~..), (~*.), enumPC, )
import Visualize (
Vis, Color(Green, Red, Blue), Plot,
fig, figP, figure, title,
xLabel, yLabel, plotD, color, label,
)

type Height = Int

data Tree = Alive Height | Hit Height | Fallen
deriving (Ord,Eq,Show)

grow :: Trans Tree
grow (Alive h) = normal [Alive k | k <- [h+1..h+5]]
grow _ = error "TreeGrowth.grow: only alive trees can grow"

hit :: Trans Tree
hit (Alive h) = certainly (Hit h)
hit _ = error "TreeGrowth.hit: only alive trees can be hit"

fall :: Trans Tree
fall _ = certainly Fallen

evolve :: Trans Tree
evolve t@(Alive _) = unfoldT (enumPC [90,4,6] [grow,hit,fall]) t
evolve t           = certainly t
-- evolve t@(Alive _) = unfoldT (enum [0.9,0.04,0.06] [grow,hit,fall]) t

{- |
tree growth simulation:
-}
seed :: Tree
seed = Alive 0

-- * exact results

-- | @tree n@ : tree distribution after n generations
tree :: Int -> Tree -> Dist Tree
tree n = n *. evolve

-- | @hist n@ : history of tree distributions for n generations
hist :: Int -> Expand Tree
hist n = n *.. evolve

-- * simulation results

{- |
Since '(*.)' is overloaded for Trans and RChange,
we can run the simulation ~. directly to @n *. live@.
-}

--simTree k n = k ~. tree n
simTree :: Int -> Int -> RTrans Tree
simTree k n = (k,n) ~*. evolve

simHist :: Int -> Int -> RExpand Tree
simHist k n = (k,n) ~.. evolve

t2 :: Dist Tree
t2  = tree 2 seed

h2 :: Space Tree
h2  = hist 2 seed

sh2, st2 :: R ()
st2 = printR \$ simTree 2000 2 seed
sh2 = printR \$ simHist 2000 2 seed

-- Alternatives:
--
-- simTree k n = k ~. n *. random evolve
-- simTree k n = (k,n) ~*. evolve

-- take a trace

height :: Tree -> Int
height Fallen = 0
height (Hit h) = h
height (Alive h) = h
{--
myPlot = plotD ((5 *. evolve) (Alive 0) >>= height)

myPlot2 = figP figure{title="Tree Growth",xLabel="Height (m)",
yLabel="Probability"}
(autoColor [
plotD ((5 *. evolve) (Alive 0) >>= height)
])

--}

p1, p2, p3, p4, p5, p6 :: Vis

p1 = fig [plotD \$ normal ([1..20]::[Int])]

p2 = fig [plotD \$ mapD height (tree 5 seed)]

p3 = figP figure{title="Tree Growth",
xLabel="Height (ft)",
yLabel="Probability"}
[plotD \$ mapD height (tree 5 seed)]

p4 = figP figure{title="Tree Growth",
xLabel="Height (ft)",
yLabel="Probability"}
[heightAtTime 5, heightAtTime 10,heightAtTime 15]

heightAtTime :: Int -> Plot
heightAtTime y = plotD \$ mapD height (tree y seed)

p5 = figP figure{title="Tree Growth",
xLabel="Height (ft)",
yLabel="Probability"}
(map heightAtTime [3,5,7])

heightCurve :: (Int,Color) -> Plot
heightCurve (n,c) = (heightAtTime n){color=c,label=show n++" Years"}

p6 = figP figure{title="Tree Growth",
xLabel="Height (ft)",
yLabel="Probability"}
(map heightCurve
[(3,Blue),(5,Green),(7,Red)])

done :: Tree -> Bool
done (Alive x) = x >= 5
done _ = True

ev5 :: Tree -> Dist Tree
ev5 = Probability.until done evolve
```