module Goal.Probability
    ( module System.Random.MWC
    , module System.Random.MWC.Monad
    , module Goal.Probability.Statistical
    , module Goal.Probability.ExponentialFamily
    , module Goal.Probability.Distributions
    , module Goal.Probability.Graphical
    , module Goal.Probability.Graphical.Harmonium
    , module Goal.Probability.Graphical.NeuralNetwork
    , module Goal.Probability
    ) where


--- Imports ---


-- Re-exports --

import System.Random.MWC hiding (uniform,uniformR)
import System.Random.MWC.Monad hiding (save)

import qualified System.Random.MWC.Monad as S (save)

import Goal.Probability.Statistical
import Goal.Probability.ExponentialFamily
import Goal.Probability.Distributions
import Goal.Probability.Graphical
import Goal.Probability.Graphical.Harmonium
import Goal.Probability.Graphical.NeuralNetwork

-- Package --

import Goal.Core
import Goal.Geometry

--- Stochastic Functions ---


seed :: RandST s Seed
-- | This little guy creates a seed. It's necessary to avoid name space
-- collisions.
seed = S.save

randomElement :: [x] -> RandST r x
-- | Returns a random element from a list.
randomElement xs = do
    u <- uniform
    let elm = round $ fromIntegral (length xs - 1) * (u :: Double)
    return $ xs !! elm

noisyFunction :: (Generative c m, Num (Sample m))
    => (c :#: m) -- ^ Noise model
    -> (x -> Sample m) -- ^ Function
    -> x
    -> RandST r (Sample m)
-- | Returns a sample from the given function with added noise.
noisyFunction m f x = do
    ns <- generate m
    return $ f x + ns

noisyRange
    :: Double -- ^ The min of the function input
    -> Double -- ^ The max function input
    -> Int -- ^ Number of samples to draw from the function
    -> Double -- ^ Standard deviation of the noise
    -> (Double -> Double) -- ^ Mixture function
    -> RandST s [(Double,Double)]
{-| Returns a set of samples from the given function with additive Gaussian noise. -}
noisyRange mn mx n sd f = do
    let xs = range mn mx n
        d = chart Standard $ fromList Normal [0,sd^2]
    fxs <- mapM (\x -> (+ f x) <$> generate d) xs
    return $ zip xs fxs