{-# LANGUAGE ExistentialQuantification #-}

module Stochastic.Distributions.Discrete(
  mkBinomial
  ,mkBernoulli
  ,mkPoisson
  ,mkZipF
  ,mkGeometric
  ,Stochastic.Distributions.Discrete.Dist(..)
  ,DiscreteDistribution(..)
  ) where

import Data.Maybe
import Control.Monad.State.Lazy
import Stochastic.Tools
import Stochastic.Generator(foldGenWhile)
import Stochastic.Distributions
import Stochastic.Distribution.Discrete
import Stochastic.Generator
import System.Random
import qualified Stochastic.Distributions.Continuous as C

data Dist = forall a . RandomGen a => Uniform Int Int a
          | Poisson C.Dist
          | forall a . RandomGen a => Geometric Double a
          | forall a . RandomGen a => Bernoulli Double a
          | forall a . RandomGen a => Binomial Int Double DiscreteCache a
          | forall a . RandomGen a => ZipF Int Double DiscreteCache a

instance RandomGen Dist where
  next g = rand g

mkBinomial :: forall a . RandomGen a => a -> Double -> Int -> Dist
mkBinomial base p n = Binomial n p cache base
  where
    nd = toDbl n
    cache :: DiscreteCache
    cache = foldr (\(w,x,y,z) l -> (w,y,z):l) [] (create n)
    create :: Int -> [(Int, Double, Double, Double)]
    create 0 = let pmfk = (nd * (log (1 - p)) )
               in [(0, pmfk, exp pmfk, exp pmfk)]
    create k = (k, lpmfk, pmfk, cdfk) : sub
      where
        sub = create (k-1)
        (j, lpmfj, pmfj, cdfj) = head sub
        pmfk = exp lpmfk
        kd = toDbl k
        lpmfk = lpmfj +
                (log p) - (log (1- p)) -
                (log (kd)) + (log (nd-kd+1))
        cdfk = pmfk + cdfj
        
-- k, pmf k, cdf k
-- do in log space
type DiscreteCache = [(Int, Double, Double)]

mkUniform :: forall a . RandomGen a => a -> Int -> Int -> Dist
mkUniform base a b = Uniform a b base

mkBernoulli :: forall a . RandomGen a => a -> Double -> Dist
mkBernoulli base p = Bernoulli p base

mkPoisson :: forall a . RandomGen a => a -> Double -> Dist
mkPoisson base y = Poisson (C.mkExp base y)

mkGeometric :: forall a . RandomGen a => a -> Double -> Dist
mkGeometric base p = Geometric p base

mkZipF :: forall a . RandomGen a => a -> Int -> Double -> Dist
mkZipF base n slope = ZipF n slope cache base
  where
    hns = sum $ take n $ harmonics slope
    cache :: DiscreteCache
    cache = create n
    create :: Int -> DiscreteCache
    create 1 = let pmfk = 1 / hns in [(1,pmfk,pmfk)]
    create k =
      (k, (1/(((toDbl k)**slope) * hns)), cdfj + pmfj) : sub
      where
        sub             = create (k-1)
        (j, pmfj, cdfj) = head sub

nextN :: forall a . RandomGen a => Int -> a -> ([Int], a)
nextN n = genTake (next) n

instance DiscreteDistribution Dist where
  rand (Binomial n p cache g0) =
    mapTuple
    (\u -> 
      length $ filter (pred u) cache  )
    (Binomial n p cache)
    (random g0)
    where pred u (k, pmf, cdf) = cdf < u
  rand (Geometric p g0) =
    mapTuple
    ((\u -> ceiling $ (log u) / (log (1-p))))
    (Geometric p)
    (random g0)
  rand (Poisson g0@(C.Exponential y u)) =
    mapTuple
    (\x -> length x)
    (Poisson)
    ((foldGenWhile (C.rand) (+) (0.0) (<1.0)) g0)
  rand (Bernoulli p g0) =
    mapTuple
    (\x -> if (x >= p) then 1 else 0) 
    (Bernoulli p)
    (random g0)
  rand (ZipF n slope cache u0) = 
    mapTuple
    (\u ->
      length $ filter (pred u) cache) 
    (ZipF n slope cache)
    (random u0)
    where pred u (k, pmf, cdf) = cdf < u
  rand (Uniform a b g0) =
    mapTuple
    (\x -> truncate (toDbl (b - a) * x + toDbl a))
    (Uniform a b)
    (random g0)

  cdf (Poisson (C.Exponential y _)) x =
    (1/(exp y)) * (sum [ (y ** (toDbl i)) / (fromInteger $ fac i) | i <- [0..x]])
  cdf (Geometric p _) x = 1 - (1-p)^x
  cdf (Bernoulli p _) x
    | x < 0     = 0
    | x >= 1    = 1
    | otherwise = p
  cdf (Binomial n p cache _) x = r
    where
      (_, _, r) = fromMaybe (0, 0, 1) $ maybeHead $ filter (\(w,_,_) -> (w==x)) cache
  cdf (ZipF n s cache _) x = r
    where
      (_, _, r) =
        fromMaybe (0, 0, 1) $ maybeHead $ filter (\(k, _, _) -> x == k) cache 
  cdf (Uniform a b _) x = toDbl (x-a) / toDbl (b-a)

  cdf' g@(Poisson (C.Exponential y _)) x =
    (sum . fst) (fold [1..])
    where
      reduce :: Double -> Int -> Double
      reduce = (\p y -> (pmf g y) + p)
      fold :: [Int] -> ([Int], [Int])
      fold = foldGenWhile (myUncons) (reduce) 0 (<x)
  cdf' (Geometric p _) x =
    ceiling $ (log (1-x)) / (log (1-p))
  cdf' (Bernoulli p _) x
    | x > p     = 1
    | otherwise = 0
  cdf' (Binomial n p cache _) x = r
    where
      (r, _, _) =
        head $ reverse $ filter (\(_, _, z) -> z > x) cache
  cdf' (ZipF n s cache _) x = r
    where
      (r, _, _) =
        head $ reverse $ filter (\(_, _, z) -> z > x) cache
  cdf' (Uniform a b _) x = truncate $ toDbl (b-a) * x + toDbl a
  
  pmf (Poisson (C.Exponential y _)) x =
    (y^x) / ( exp y * (fromInteger $ fac x) )
  pmf (Geometric p _) x = 1 - (1-p)^(x-1)
  pmf (Bernoulli p _) 0 = 1-p
  pmf (Bernoulli p _) 1 = p
  pmf (Binomial n p cache _) x = r
    where
      (_, r, _) = head $ filter (\(w,_,_) -> (w==x)) cache
  pmf (ZipF n s cache _) k = r
    where
      (_, r, _) = head $ filter (\(w,_,_) -> (w==k)) cache
  pmf (Uniform a b _) x = toDbl x/toDbl (b-a)


myUncons :: [a] -> (a, [a])
myUncons (x:xs) = (x, xs)

toDbl = fromInteger . toInteger