{-| A general purpose library for simulating differential processes, of a deterministic or
    stochastic nature. -}

module Goal.Simulation.Mealy
    ( -- * Exports
      module Data.Machine
    -- * Accumulation
    , accumulateFunction
    , accumulateFunction'
    , accumulateMealy
    , accumulateMealy'
    , accumulateRandomFunction
    , accumulateRandomFunction'
    , accumulateRandomFunction0
      -- * Execution
    , stream
    , streamM
    , streamM_
    ) where

--- Imports ---


-- Goal --

import Goal.Probability

-- Reexporting --

import Data.Machine

import qualified Control.Monad.ST as ST


--- Mealys ---


accumulateFunction :: (a -> acc -> (b,acc)) -> acc -> Mealy a b
-- | accumulateFunction takes a function from a value and an accumulator (e.g. just a sum
-- value or an evolving set of parameters for some model) to a value and an accumulator.
-- The accumulator is then looped back into the function, returning a Mealy from a to
-- b, which updates the accumulator every time step.
accumulateFunction f acc = Mealy $ \a ->
    let (b,acc') = f a acc
     in (b,accumulateFunction f acc')

accumulateFunction' :: (a -> acc -> (b,acc)) -> acc -> Mealy a (b,acc)
-- | accumulateFunction' acts like accumulateFunction but the Mealy automata will
-- continue to return the accumulator as it generates it.
accumulateFunction' f =
    accumulateFunction f'
    where f' a acc =
              let (b,acc') = f a acc
               in ((b,acc'),acc')

accumulateRandomFunction :: (a -> acc -> forall s . RandST s (b,acc)) -> acc -> RandST s' (Mealy a b)
-- | accumulateRandomFunction is analogous to accumulateFunction, but takes as an
-- argument a function which returns a random variable.
accumulateRandomFunction rf acc0 = do
    rf' <-  accumulateRandomFunction0 (uncurry rf)
    return $ accumulateMealy acc0 rf'

accumulateRandomFunction' :: (a -> acc -> forall s . RandST s (b,acc)) -> acc -> RandST s' (Mealy a (b,acc))
-- | accumulateRandomFunction' is analogous to accumulateFunction', but takes as an
-- argument a function which returns a random variable.
accumulateRandomFunction' rf acc0 = do
    rf' <- accumulateRandomFunction0 (uncurry rf)
    return $ accumulateMealy' acc0 rf'

accumulateRandomFunction0 :: (a -> forall s . RandST s b) -> RandST s' (Mealy a b)
-- | accumulateRandomFunction' Mealifies stateless random functions.
accumulateRandomFunction0 rf = do
    sd <- seed
    return $ accumulateFunction f sd
    where f a sd = ST.runST $ do
              gn <- restore sd
              b <- runRand (rf a) gn
              sd' <- save gn
              return (b,sd')

accumulateMealy :: acc -> Mealy (a,acc) (b,acc) -> Mealy a b
-- | accumulateMealy takes a Mealy with an accumulating parameter and loops it.
accumulateMealy acc0 mly0 =
    accumulateFunction f (acc0,mly0)
    where f a (acc,Mealy cf) =
              let ((b,acc'),mly') = cf (a,acc)
               in (b,(acc',mly'))

accumulateMealy' :: acc -> Mealy (a,acc) (b,acc) -> Mealy a (b,acc)
-- | accumulateMealy except with a returned accumulator.
accumulateMealy' acc0 mly0 =
    accumulateFunction f (acc0,mly0)
    where f a (acc,Mealy cf) =
              let ((b,acc'),mly') = cf (a,acc)
               in ((b,acc'),(acc',mly'))

--- Execution ---

{-
parallelizeMealys :: [Mealy a b] -> Mealy [a] [b]
{-| Turns a list of circuits into a circuit over lists, bound by the power of parMap rseq. -}
parallelizeMealys crcs = Mealy $ \as ->
    let (bs,crcs') = unzip $ parZip crcs as
    in  (bs, parallelizeMealys crcs')
    where parZip [] [] = []
          parZip (crc:crcs) (a:as) =
              let (b,crc') = runMealy crc a
              in  b `par` (b,crc') : parZip crcs as
          parZip _ _ = error "Parallel circuit does not match size of input"
          -}

stream :: Mealy a b -> [a] -> [b]
stream mly as = run . supply as . auto $ mly

streamM :: Monad m => Mealy a b -> (b -> m c) -> [a] -> m [c]
streamM mly fM as = runT . supply as $ auto mly ~> autoM fM

streamM_ :: Monad m => Mealy a b -> (b -> m c) -> [a] -> m ()
streamM_ mly fM as = runT_ . supply as $ auto mly ~> autoM fM