{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE NoMonomorphismRestriction #-}

{-| Simple moving average style folds -}

module Control.Foldl.Incremental.Simple (
    -- * incrementalize
    incrementalizeSimple
    -- * common simple folds
  , sma
  , sabsma
  , ssqma
  , sstd
  ) where

import           Control.Applicative
import           Control.Foldl as L
import           Data.Sequence (ViewR(EmptyR, (:>)), (<|))
import qualified Data.Sequence as Seq

(</>) :: (Fractional c, Applicative f) => f c -> f c -> f c
(</>) = liftA2 (/) 

(<+>) :: (Fractional c, Applicative f) => f c -> f c -> f c
(<+>) = liftA2 (+) 

maybeSum :: Fold (Maybe Double) (Maybe Double)
maybeSum = Fold (<+>) (Just 0) id

{-|
Incrementalize takes a function and turns it into a `Control.Foldl.Fold` where the step is an Increment similar to the typical step in a simple moving average calculation.
-}
incrementalizeSimple :: (a -> Double) -> Int -> Fold a Double
incrementalizeSimple f n = L.Fold step begin done
  where
    begin = Seq.replicate n Nothing
    av x = L.fold maybeSum x </> pure (fromIntegral n)
    done x = case av x of
        Nothing -> 0/0
        Just x' -> x'
    step x a = Just (f a) <| pop x
    pop x = case Seq.viewr x of
        EmptyR   -> x
        x'' :> _ -> x''

-- | a simple moving average
sma :: Int -> L.Fold Double Double
sma = incrementalizeSimple id

-- | simple squared moving average
ssqma :: Int -> L.Fold Double Double
ssqma = incrementalizeSimple (**2)

-- | simple absolute moving average
sabsma :: Int -> L.Fold Double Double
sabsma = incrementalizeSimple abs

-- | simple standard deviation
sstd :: Int -> Fold Double Double
sstd n = (\s ss -> sqrt (ss - s**2)) <$> sma n <*> ssqma n