
{-# LANGUAGE FunctionalDependencies #-}

module Distribution ( 
    Dist, pdf, logpdf, samplefrom,
    Uniform(..), Normal(..), Bernoulli(..), 
    Exponential(..), Gamma(..), Poisson(..), 
    Binomial(..), Beta(..), 
    Prob, Weight, Weighted
) where

import Randomness ( Random(..), uniform, Prob, Weight, Weighted )
import Numeric.Log ( Log(..) )
import qualified Data.Random.Distribution.Uniform as R ( uniform )
import qualified Data.Random.Distribution.Normal as R ( normal )
import qualified Data.Random.Distribution.Bernoulli as R ( bernoulli )
import qualified Data.Random.Distribution.Binomial as R ( binomial )
import qualified Data.Random.Distribution.Poisson as R ( poisson )
import qualified Statistics.Distribution as S ( ContDistr(density, logDensity, quantile), DiscreteDistr(probability, logProbability) )
import qualified Statistics.Distribution.Uniform as S ( uniformDistr )
import qualified Statistics.Distribution.Normal as S ( normalDistr )
import qualified Statistics.Distribution.Binomial as S ( binomial )
import qualified Statistics.Distribution.Poisson as S ( poisson )
import qualified Statistics.Distribution.Gamma as S ( gammaDistr )
import qualified Statistics.Distribution.Exponential as S ( exponential )
import qualified Statistics.Distribution.Binomial as S ( binomial )
import qualified Statistics.Distribution.Beta as S ( betaDistr )
import qualified Statistics.Distribution.DiscreteUniform as S ( discreteUniformAB )

class Dist d a | d -> a where
    samplefrom :: d -> Random a

    pdf :: d -> a -> Double
    pdf d = exp . ln . logpdf d

    logpdf :: d -> a -> Weight
    logpdf d = Exp . log . pdf d


data Uniform = Uniform Double Double
-- data DiscreteUniform = DiscreteUniform Int Int
data Normal = Normal Double Double
data Bernoulli = Bernoulli Prob
data Exponential = Exponential Double
data Gamma = Gamma Double Double
data Poisson = Poisson Double
data Binomial = Binomial Int Prob
data Beta = Beta Double Double

instance Dist Uniform Double where
    pdf (Uniform a b) = S.density $ S.uniformDistr a b
    logpdf (Uniform a b) = Exp . S.logDensity (S.uniformDistr a b)
    samplefrom (Uniform a b) = do 
        u <- uniform
        return $ a + u * (b - a)

{-
instance Dist DiscreteUniform Int where
    pdf (DiscreteUniform a b) = S.probability $ S.discreteUniformAB a b
    logpdf (DiscreteUniform a b) = Exp . S.logProbability (S.discreteUniformAB a b)
    samplefrom (DiscreteUniform a b) = uniform (a, b)
-}

instance Dist Normal Double where
    pdf (Normal mu sig) = S.density $ S.normalDistr mu sig
    logpdf (Normal mu sig) = Exp . S.logDensity (S.normalDistr mu sig)
    samplefrom (Normal mu sig) = S.quantile (S.normalDistr mu sig) <$> uniform

instance Dist Bernoulli Bool where
    pdf (Bernoulli p) b = if b then p else 1-p
    logpdf (Bernoulli p) b = if b then Exp $ log p else Exp $ log (1-p)
    samplefrom (Bernoulli p) = (<= p) <$> uniform

instance Dist Gamma Double where
    pdf (Gamma k theta) = S.density $ S.gammaDistr k theta
    logpdf (Gamma k theta) = Exp . S.logDensity (S.gammaDistr k theta)
    samplefrom (Gamma k theta) = S.quantile (S.gammaDistr k theta) <$> uniform

instance Dist Exponential Double where
    pdf (Exponential lambda) = S.density $ S.exponential lambda
    logpdf (Exponential lambda) = Exp . S.logDensity (S.exponential lambda)
    samplefrom (Exponential lambda) = S.quantile (S.exponential lambda) <$> uniform

instance Dist Beta Double where
    pdf (Beta a b) = S.density $ S.betaDistr a b
    logpdf (Beta a b) = Exp . S.logDensity (S.betaDistr a b)
    samplefrom (Beta a b) = do
        x <- samplefrom (Gamma a 1)
        y <- samplefrom (Gamma b 1)
        return $ x / (x+y)

instance Dist Poisson Int where
    pdf (Poisson lambda) = S.probability $ S.poisson lambda
    logpdf (Poisson lambda) = Exp . S.logProbability (S.poisson lambda)
    samplefrom (Poisson lambda) = step 0 1.0  
        where 
            l = exp (-lambda)
            step k p = do
                u <- uniform
                let p' = p*u
                if p' <= l then return k 
                           else step (k+1) p'


instance Dist Binomial Int where
    pdf (Binomial n p) = S.probability $ S.binomial n p
    logpdf (Binomial n p) = Exp . S.logProbability (S.binomial n p)
    samplefrom d@(Binomial n p) | n <= 10 = sumBernoullis n 0
                                | otherwise = do
        let a = 1 + n `div` 2
        let b = 1 + n - a
        x <- samplefrom $ Beta (fromIntegral a) (fromIntegral b) 
        if x >= p 
            then samplefrom $ Binomial (a-1) (p/x)
            else do 
                s <- samplefrom (Binomial (b-1) ((p-x)/(1-x)))
                return (s + a)
        where 
            sumBernoullis 0 s = return s
            sumBernoullis n s = do
                u <- uniform
                if u <= p then sumBernoullis (n-1) (s+1) else sumBernoullis (n-1) s
