{-# LANGUAGE ConstraintKinds        #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE GADTs                  #-}
{-# LANGUAGE InstanceSigs           #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE QuantifiedConstraints  #-}
{-# LANGUAGE RankNTypes             #-}
{-# LANGUAGE TemplateHaskell        #-}
{-# LANGUAGE TypeOperators          #-}
{-# LANGUAGE UndecidableInstances   #-}

module Q.Stats.Arima where
import           Control.Monad.State
import           Data.Foldable
import           Data.Functor.Identity
import           Data.Random
import           Data.Random.Source
import           Data.RVar
import           Data.Time
import           Numeric.LinearAlgebra
import           Q.Stats.TimeSeries
import           System.Random.Mersenne.Pure64
import           Data.Random.Distribution
import           Data.Random.Distribution.Poisson
import           Data.Random.Distribution.T
import           Data.RVar
import           Statistics.Sample

data Ewma d = Ewma Double d

--ll :: (Ewma d) -> [DataPoint Double] -> (Double -> Double)
ll :: Ewma (d Double)
-> t (DataPoint LocalTime Double)
-> StateT Double Identity (t Double)
ll (Ewma Double
lambda d Double
d) t (DataPoint LocalTime Double)
datapoints = (DataPoint LocalTime Double -> StateT Double Identity Double)
-> t (DataPoint LocalTime Double)
-> StateT Double Identity (t Double)
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
mapM DataPoint LocalTime Double -> StateT Double Identity Double
ll_ t (DataPoint LocalTime Double)
datapoints where
  ll_ :: DataPoint LocalTime Double -> State Double Double
  ll_ :: DataPoint LocalTime Double -> StateT Double Identity Double
ll_ x :: DataPoint LocalTime Double
x@(DataPoint LocalTime
_ Double
v) = do
    Double
vart <- StateT Double Identity Double
forall s (m :: * -> *). MonadState s m => m s
get
    let vart2 :: Double
vart2 = Double
lambda Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
vart Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lambda) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v
    Double -> StateT Double Identity ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put Double
vart2
    Double -> StateT Double Identity Double
forall (m :: * -> *) a. Monad m => a -> m a
return (Double -> StateT Double Identity Double)
-> Double -> StateT Double Identity Double
forall a b. (a -> b) -> a -> b
$ d Double -> Double -> Double
forall (d :: * -> *) t. PDF d t => d t -> t -> Double
logPdf d Double
d (Double -> Double
forall a. Floating a => a -> a
sqrt (Double
v  Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
v Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
vart))


--forecast :: (Distribution d Double) => (Ewma d) -> Int ->
forecast :: forall d. (Distribution d Double) => Ewma (d Double) -> StateT Double RVar Double
forecast :: Ewma (d Double) -> StateT Double RVar Double
forecast (Ewma Double
lambda d Double
d) = do
  Double
y <- RVarT Identity Double -> StateT Double RVar Double
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (RVarT Identity Double -> StateT Double RVar Double)
-> RVarT Identity Double -> StateT Double RVar Double
forall a b. (a -> b) -> a -> b
$ d Double -> RVarT Identity Double
forall (d :: * -> *) t. Distribution d t => d t -> RVar t
rvar d Double
d
  Double
vart <- StateT Double RVar Double
forall s (m :: * -> *). MonadState s m => m s
get
  let vart2 :: Double
vart2 = Double
lambda Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
vart Double -> Double -> Double
forall a. Num a => a -> a -> a
+ (Double
1 Double -> Double -> Double
forall a. Num a => a -> a -> a
- Double
lambda) Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double
y
  Double -> StateT Double RVar ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put Double
vart2
  Double -> StateT Double RVar Double
forall (m :: * -> *) a. Monad m => a -> m a
return (Double
y Double -> Double -> Double
forall a. Num a => a -> a -> a
* Double -> Double
forall a. Floating a => a -> a
sqrt Double
vart)


--forecastN :: Distribution d Double => Ewma (d Double) -> Int -> Double -> RVar ([Double], Double)
forecastN :: Ewma (d Double) -> Double -> Int -> m ([Double], Double)
forecastN Ewma (d Double)
ewma Double
var0 Int
n =  RVar ([Double], Double) -> m ([Double], Double)
forall (d :: * -> *) (m :: * -> *) t.
(Sampleable d m t, MonadRandom m) =>
d t -> m t
sample (RVar ([Double], Double) -> m ([Double], Double))
-> RVar ([Double], Double) -> m ([Double], Double)
forall a b. (a -> b) -> a -> b
$ StateT Double RVar [Double] -> Double -> RVar ([Double], Double)
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
runStateT (Int -> StateT Double RVar Double -> StateT Double RVar [Double]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
n (Ewma (d Double) -> StateT Double RVar Double
forall (d :: * -> *).
Distribution d Double =>
Ewma (d Double) -> StateT Double RVar Double
forecast Ewma (d Double)
ewma)) Double
var0