{-# LANGUAGE TypeFamilies #-}
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(Generator(..), foldWhile)
import Stochastic.Distributions
import Stochastic.Distribution.Discrete
import qualified Stochastic.Distributions.Continuous as C

instance Generator Dist where
  type (From Dist) = Int
  nextG = state $ \ g0 -> rand g0

data Dist =
  Uniform Int Int UniformBase
  | Poisson C.Dist
  | Geometric Double UniformBase
  | Bernoulli Double UniformBase
  | Binomial Int Double DiscreteCache UniformBase
  | ZipF Int Double DiscreteCache UniformBase

mkBinomial :: UniformBase -> 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 :: UniformBase -> Int -> Int -> Dist
mkUniform base a b = Uniform a b base

mkBernoulli :: UniformBase -> Double -> Dist
mkBernoulli base p = Bernoulli p base

mkPoisson :: UniformBase -> Double -> Dist
mkPoisson base y = Poisson (C.mkExp base y)

mkGeometric :: UniformBase -> Double -> Dist
mkGeometric base p = Geometric p base

mkZipF :: UniformBase -> 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

instance DiscreteDistribution Dist where
  rand (Binomial n p cache g0) =
    mapTuple
    (\u -> 
      length $ filter (pred u) cache)
    (Binomial n p cache)
    (C.rand g0)
    where pred u (k, pmf, cdf) = cdf < u
  rand (Geometric p g0) =
    mapTuple
    (\u -> ceiling $ (log u) / (log (1-p)))
    (Geometric p)
    (C.rand g0)
  rand (Poisson g0@(C.Exponential y u)) =
    mapTuple
    (\x -> length x)
    (Poisson)
    (runState (foldWhile (+) 0.0 (<1.0)) g0)
  rand (Bernoulli p g0) =
    mapTuple
    (\x -> if (x >= p) then 1 else 0)
    (Bernoulli p)
    (C.rand g0)
  rand (ZipF n slope cache u0) = 
    mapTuple
    (\u ->
      length $ filter (pred u) cache)
    (ZipF n slope cache)
    (C.rand 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)
    (C.rand 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) = head $ filter (\(w,_,_) -> (w==x)) cache
  cdf (ZipF n s cache _) x = r
    where
      (_, _, r) =
        head $ 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 $ runState (fold) [1..]
    where
      reduce :: Int -> Double -> Double
      reduce = (\y p -> (pmf g y) + p)
      fold :: State [Int] [From [Int]]
      fold = foldWhile (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)


toDbl = fromInteger . toInteger