module Goal.Probability.ExponentialFamily (
    -- * Exponential Families
    ExponentialFamily (sufficientStatistic, baseMeasure)
    , sufficientStatisticN
    -- ** Dual Parameters
    , Natural (Natural)
    , Mixture (Mixture)
    -- ** Divergence
    , klDivergence
    , relativeEntropy
    ) where

--- Imports ---


-- Package --

import Goal.Probability.Statistical

import Goal.Geometry


--- Exponential Families ---


-- | A 'Statistical' 'Manifold' is a member of the 'ExponentialFamily' if we can
-- specify a 'sufficientStatistic' of fixed length. Defining the 'baseMeasure'
-- is also necessary in order to render unique the 'Natural' and 'Mixture'
-- parameterizations.
--
-- 'ExponentialFamily' distributions theoretically have a 'Riemannian' geometry
-- given by the Fisher information metric, given rise to the 'DualChart' system
-- of 'Natural' and 'Mixture'. A 'Point' on the 'ExponentialFamily' 'Manifold' in
-- one of these dual coordinates is assumed to be equipped the corresponding
-- dual connection. Under this assumption, we take the 'Manifold' itself to be
-- self-dual to simplify types.
class (Statistical m, Legendre Natural m, Legendre Mixture m) => ExponentialFamily m where
    sufficientStatistic :: m -> Sample m -> Mixture :#: m
    baseMeasure :: m -> Sample m -> Double

sufficientStatisticN :: ExponentialFamily m => m -> [Sample m] -> Mixture :#: m
-- | The sufficient statistic of N iid random variables.
sufficientStatisticN m xs =
    fromIntegral (length xs) /> foldr1 (<+>) (sufficientStatistic m <$> xs)

klDivergence
    :: (ExponentialFamily m, Transition c Natural m, Transition d Mixture m)
    => c :#: m -> d :#: m -> Double
klDivergence q p = divergence (chart Natural $ transition q) (chart Mixture $ transition p)

relativeEntropy
    :: (ExponentialFamily m, Transition c Mixture m, Transition d Natural m)
    => c :#: m -> d :#: m -> Double
relativeEntropy p q = klDivergence q p

-- | A parameterization in terms of the natural coordinates of an exponential family.
data Natural = Natural

-- | A representation in terms of the mean sufficient statistics of an exponential family.
data Mixture = Mixture

instance Primal Natural where
    type Dual Natural = Mixture

instance Primal Mixture where
    type Dual Mixture = Natural


--- Instances ---


-- Generic --

instance ExponentialFamily m => MaximumLikelihood Mixture m where
    mle = sufficientStatisticN

instance ExponentialFamily m => MaximumLikelihood Natural m where
    mle m xs = potentialMapping $ sufficientStatisticN m xs

-- Replicated --

instance ExponentialFamily m => ExponentialFamily (Replicated m) where
    sufficientStatistic (Replicated m _) xs =
        joinReplicated $ sufficientStatistic m <$> xs
    baseMeasure (Replicated m _) xs = product $ baseMeasure m <$> xs

-- Fisher Manifolds --

instance ExponentialFamily m => AbsolutelyContinuous Natural m where
    density p x =
        let s = manifold p
         in exp ((p <.> sufficientStatistic s x) - potential p) * baseMeasure s x

instance ExponentialFamily m => Transition Mixture Natural m where
    transition = potentialMapping

instance ExponentialFamily m => Transition Natural Mixture m where
    transition = potentialMapping