{-# LANGUAGE BangPatterns #-}
module QuantLib.Stochastic.Process
        ( module QuantLib.Stochastic.Process )
        where

import           Control.Monad              (foldM)
import           Data.List                  (foldl')
import           QuantLib.Stochastic.Random (NormalGenerator (..))

-- | Discretization of stochastic process over given interval
class Discretize b where
        dDrift :: StochasticProcess a => a->b->Dot->Double
        dDiff  :: StochasticProcess a => a->b->Dot->Double
        dDt    :: StochasticProcess a => a->b->Dot->Double

-- | 1D Stochastic process
class StochasticProcess a where
        drift  :: a->Dot->Double
        diff   :: a->Dot->Double
        evolve :: Discretize b=> b->a->Dot->Double->Dot
        evolve discr p dot dw = Dot newT newX
                where   !newT = getT dot + dDt p discr dot
                        !newX = getX dot + dDrift p discr dot + dDiff p discr dot * dw

-- | Dot. t and x pair
data Dot = Dot { getT :: {-# UNPACK #-} !Double, getX :: {-# UNPACK #-} !Double }
        deriving (Show, Eq)

-- | Path as list of Dots
type Path = [Dot]

-- | Generates sample path for given stochastic process under discretization and normal generator for given amount of steps, starting from x0
generatePath :: (StochasticProcess a, NormalGenerator b, Discretize c) => b->c->a->Int->Dot->IO Path
generatePath rnd discr sp steps x0 = do
        (!list, _) <- foldM generator ([], rnd) [1..steps]
        let !path = foldl' evolver [x0] list
        return $! reverse path
        where   evolver p dw = evolve discr sp (head p) dw : p
                generator (list, r) _ = do
                        (!p, newRnd) <- ngGetNext r
                        return (p:list, newRnd)


-- | Geometric Brownian motion
data GeometricBrownian = GeometricBrownian {
        gbDrift :: Double,
        gbDiff  :: Double
        } deriving (Show)

instance StochasticProcess GeometricBrownian where
        drift p (Dot _ x) = gbDrift p * x
        diff  p (Dot _ x) = gbDiff p  * x

-- | Ito process
data ItoProcess = ItoProcess {
        ipDrift :: Dot->Double,
        ipDiff  :: Dot->Double
        }

instance StochasticProcess ItoProcess where
        drift   = ipDrift
        diff    = ipDiff

-- | Square-root process
data SquareRootProcess = SquareRootProcess {
        srpSpeed :: Double,
        srpMean  :: Double,
        srpSigma :: Double
        } deriving (Show)

instance StochasticProcess SquareRootProcess where
       drift p (Dot _ x) = srpSpeed p * (srpMean p - x)
       diff  p (Dot _ x) = srpSigma p * sqrt x

-- | Ornstein-Uhlenbeck process
data OrnsteinUhlenbeckProcess = OrnsteinUhlenbeckProcess {
        oupSpeed :: Double,
        oupLevel :: Double,
        oupSigma :: Double
        } deriving (Show)

instance StochasticProcess OrnsteinUhlenbeckProcess where
        drift p (Dot _ x) = oupSpeed p * (oupLevel p - x)
        diff  p _         = oupSigma p

-- | Generalized Black-Scholes process
data BlackScholesProcess = BlackScholesProcess {
        bspRiskFree :: Double->Double,
        bspDividend :: Double->Double,
        bspBlackVol :: Dot->Double
        }

instance StochasticProcess BlackScholesProcess where
        drift (BlackScholesProcess r q v) dot = r (getT dot) - q ( getT dot) - 0.5 * v dot ** 2
        diff = bspBlackVol