module Goal.Probability.Statistical (
    -- * Stastical Manifolds
      Statistical (sampleSpace)
    , Sample
    , samples
    , SampleSpace
    -- ** Standard Chart
    , Standard (Standard)
    , standardGenerate
    -- ** Distributions
    , Generative (generate)
    , AbsolutelyContinuous (density)
    , expectation
    , MaximumLikelihood (mle)
    ) where


--- Imports ---


-- Package --

import Goal.Geometry

-- Unqualified --

import System.Random.MWC.Monad


--- Test Bed ---


--- Probability Theory ---


-- | A 'Statistical' 'Manifold' is a 'Manifold' of probability distributions,
-- which all have in common a particular 'SampleSpace'.
class (Set (SampleSpace m), Manifold m) => Statistical m where
    type SampleSpace m :: *
    sampleSpace :: m -> SampleSpace m

-- | A 'Sample' is an 'Element' of the 'SampleSpace'.
type Sample m = Element (SampleSpace m)

samples :: (Discrete (SampleSpace m), Statistical m) => m -> [Sample m]
-- | The list of 'Sample's.
samples = elements . sampleSpace

-- | A distribution is 'Generative' if we can 'generate' samples from it. Generation is
-- powered by MWC Monad.
class Statistical m => Generative c m where
    generate :: c :#: m -> RandST r (Sample m)

-- | If a distribution is 'AbsolutelyContinuous' with respect to a reference
-- measure on its 'SampleSpace', then we may define the 'density' of a
-- probability distribution as the Radon-Nikodym derivative of the probability
-- measure with respect to the base measure.
class Statistical m => AbsolutelyContinuous c m where
    density :: c :#: m -> Sample m -> Double

-- | 'expectation' computes the brute force expected value of a 'Discrete' set given an appropriate 'density'.
expectation :: (AbsolutelyContinuous c m, Discrete (SampleSpace m)) => c :#: m -> (Sample m -> Double) -> Double
expectation p f =
    let xs = elements . sampleSpace $ manifold p
     in sum $ zipWith (*) (f <$> xs) (density p <$> xs)


-- | 'mle' computes the 'MaximumLikelihood' estimator.
class Statistical m => MaximumLikelihood c m where
    mle :: m -> [Sample m] -> c :#: m

-- Standard Chart --

-- | A parameterization which represents the standard or typical parameterization of
-- the given manifold, e.g. the 'Poisson' rate or 'Normal' mean and standard deviation.
data Standard = Standard deriving (Eq, Read, Show)

standardGenerate :: (Manifold m, Generative Standard m, Transition c Standard m) => c :#: m -> RandST r (Sample m)
standardGenerate = generate . chart Standard . transition

--- Instances ---


-- DirectSums --

instance (Statistical m, Statistical n) => Statistical (m,n) where
    type SampleSpace (m,n) = (SampleSpace m, SampleSpace n)
    sampleSpace (m,n) = (sampleSpace m,sampleSpace n)

instance (Generative c m, Generative c n) => Generative c (m,n) where
    generate cmn = do
        let (cm,cn) = splitPair' cmn
        mx <- generate cm
        nx <- generate cn
        return (mx, nx)

instance (AbsolutelyContinuous Standard m, AbsolutelyContinuous Standard n) => AbsolutelyContinuous Standard (m,n) where
    density cmn (mx,nx) =
        let (cm,cn) = splitPair' cmn
        in density cm mx * density cn nx

-- Replicated --

instance Statistical m => Statistical (Replicated m) where
    type SampleSpace (Replicated m) = Replicated (SampleSpace m)
    sampleSpace (Replicated m n) = Replicated (sampleSpace m) n

instance (Statistical m, Generative c m) => Generative c (Replicated m) where
    generate = sequence . mapReplicated generate

instance (Statistical m, AbsolutelyContinuous Standard m) => AbsolutelyContinuous Standard (Replicated m) where
    density ds xs = product $ zipWith ($) (mapReplicated density ds) xs

instance (Statistical m, Transition Standard c m) => Transition Standard c (Replicated m) where
    transition = joinReplicated . mapReplicated transition

instance (Statistical m, Transition c Standard m) => Transition c Standard (Replicated m) where
    transition = joinReplicated . mapReplicated transition


--- Graveyard ---


{-
manifoldExpectation :: (Manifold n, AbsolutelyContinuous c m, Discrete (SampleSpace m))
    => c :#: m -> (Sample m -> d :#: n) -> d :#: n
manifoldExpectation p f =
    let xs = elements . sampleSpace $ manifold p
     in foldl1' (<+>) $ zipWith (.>) (density p <$> xs) (f <$> xs)

-}