module Stochastic.Distributions(
  ContinuousDistribution(..)
  ,DiscreteDistribution(..)
  ,Distributions(..)
  ,mkDistributions
  ,stdDistributions
  ,liftD
  ,liftC
  ,liftDN
  ,liftCN
--  ,plot
  ) where

import Helpers(histogram, statefully, mapTuple)
import Stochastic.Distribution

import qualified Stochastic.Uniform     as Uni
import qualified Stochastic.ZipF        as Zipf
import qualified Stochastic.Geometric   as Geo
import qualified Stochastic.Exponential as Exp
import qualified Stochastic.Poisson     as Poi
import qualified Stochastic.Normal      as Nor
import qualified Stochastic.Bernoulli   as Ber
import qualified Stochastic.Binomial    as Bin

data Distributions = Distributions {
  mkUniform    :: Int                     -> Uni.Uniform
  ,mkExp       :: Int -> Double           -> Exp.Exponential
  ,mkNormal    :: Int -> Double -> Double -> Nor.Normal

  ,mkZipF      :: Int -> Int    -> Double -> Zipf.ZipF
  ,mkGeometric :: Int -> Double           -> Geo.Geometric
  ,mkPoisson   :: Int -> Double           -> Poi.Poisson
  ,mkBernoulli :: Int -> Double           -> Ber.Bernoulli
  ,mkBinomial  :: Int -> Double -> Int    -> Bin.Binomial
}

stdDistributions = mkDistributions Uni.stdUniform

mkDistributions uniform = Distributions {
  mkUniform     = uniform
  , mkZipF      = apiMkZipF      (uniform)
  , mkGeometric = apiMkGeometric (uniform)
  , mkExp       = apiMkExp       (uniform)
  , mkPoisson   = apiMkPoisson   (uniform)
  , mkNormal    = apiMkNormal    (uniform)
  , mkBernoulli = apiMkBernoulli (uniform)
  , mkBinomial  = apiMkBinomial  (uniform)
}

apiMkZipF      mkUni seed n slope = Zipf.mkZipF      (mkUni seed) n slope
apiMkGeometric mkUni seed p       = Geo.mkGeometric  (mkUni seed) p
apiMkExp       mkUni seed y       = Exp.mkExp        (mkUni seed) y
apiMkPoisson   mkUni seed y       = Poi.mkPoisson $  apiMkExp (mkUni) seed y
apiMkNormal    mkUni seed m d     = Nor.mkNormal     (mkUni seed) m d
apiMkBernoulli mkUni seed d       = Ber.mkBernoulli  (mkUni seed) d
apiMkBinomial  mkUni seed d n     = Bin.mkBinomial   (apiMkBernoulli (mkUni) seed d) n

liftD :: (DiscreteDistribution g) => (Int, Int) -> (Int -> a) -> (g -> (a, g))
liftD range f g0 = ((f r), g1)
  where
    (r, g1) = (randIntIn range g0)

liftDN :: (DiscreteDistribution g) => [(Int, Int)] -> ([Int] -> a) -> (g -> (a, g))
liftDN ranges f g0 = mapTuple (f) (id) (h ranges g0)
  where
    h [] g1 = ([], g1)
    h (r:rs) g1 = let (s, g2) = randIntIn r g1 in mapTuple (\x -> s:x) (id) (h rs g2)

liftC :: (ContinuousDistribution g) => (Double -> a) -> (g -> (a, g))
liftC f g0 = ((f r), g1)
  where
    (r, g1) = (randDouble g0)

liftCN :: (ContinuousDistribution g) => Int -> ([Double] -> a) -> (g -> (a, g))
liftCN n f g0 = mapTuple (f) (id) (randDoubles n g0)

plot :: DiscreteDistribution g => g -> Int -> Int -> [Int]
plot g0 interval samples = []
{-
plot g n samples = map (truncate . (+0.5) . (*100))
  $ map
  (\x ->
    (toDbl x)/(toDbl samples))
  $ histogram (1.0 / (toDbl n))
  (fst (randInts g samples))
-}
toDbl = fromInteger . toInteger