{-# LANGUAGE ExistentialQuantification #-}

module Stochastic.Distributions.Continuous(
  mkUniform
  ,mkExp
  ,mkNormal
  ,mkEmpirical
  ,Dist(..)
  ,ContinuousDistribution(..)
  ) where

import Data.Maybe
import Control.Monad.State.Lazy
--import Stochastic.Analysis
import Stochastic.Generator

import Stochastic.Distributions(stdBase)
import qualified Stochastic.Distributions as B(cdf, mkEmpirical, Empirical)
import Stochastic.Distribution.Continuous
import Stochastic.Tools
import Data.Number.Erf
import System.Random

data Dist = forall a . RandomGen a => Uniform a
          | forall a . RandomGen a => Exponential Double a
          | forall a . RandomGen a => Normal Double Double (Maybe Double) a
          | forall a . RandomGen a => ChiSquared Int a
          | forall a . RandomGen a => Empirical B.Empirical a
    -- empirical points, lo, [(point, mass)]
instance RandomGen Dist where
  next g@(Uniform uni) = mapTuple id Uniform $ next g
  next g@(Exponential y _) =
    let (x, g') = rand g in
    let (_, scale) = genRange g in
    let x' = truncate ((fromIntegral scale) * x) in
    if (x' < scale)
    then (x', g')
    else next g'
  next g@(Normal mean dev _ _) =
    let (x, g') = rand g in
    (truncate x, g')

  genRange g@(Uniform _) = (minBound :: Int, maxBound :: Int)
  genRange g@(Exponential _ _) = (0, maxBound `div` 4096 :: Int)
  genRange g@(Normal mean dev _ _) =
    let pm = (dev * 6.66) in
    (ceiling $ mean - pm, ceiling $ mean + pm)


mkEmpirical :: forall a . RandomGen a => a -> [Double] -> Dist
mkEmpirical base samples = Empirical (B.mkEmpirical samples) base

mkExp :: forall a . RandomGen a => a ->  Double -> Dist
mkExp base y = Exponential y base

mkNormal :: forall a . RandomGen a => a -> Double -> Double -> Dist
mkNormal uni mean dev = Normal mean dev Nothing uni

mkUniform :: forall a . RandomGen a => a -> Dist
mkUniform uni = Uniform uni

intWordDbl :: Int -> Double
intWordDbl x = fromRational $ toRational ((fromInteger $ toInteger x) :: Word)

randomN :: forall a . forall b . (RandomGen a, Random b) => Int -> a -> ([b], a)
randomN n = genTake (random) n
  

instance ContinuousDistribution Dist where
  rand (Uniform uni) = mapTuple (id) (Uniform) (random uni)
  rand (Exponential y u) =
    mapTuple ((\x -> -(log $ x) / y)) (Exponential y) (random u)
  rand (Normal mean dev m uni) = f m
    where
      f (Just x) = (x, (Normal mean dev Nothing uni'))
      f Nothing  = (y, (Normal mean dev (Just z) uni'))
      (vs, uni') = randomN 2 uni
      [u1, u2] = map (id) vs
      from_u g = mean + dev * (sqrt (-2 * (log u1))) * ( g (2 * pi * u2) )
      y = from_u (sin)
      z = from_u (cos)

  cdf  (Uniform _) x = x
  cdf  (Exponential y _) x = 1 - (1 / (exp (y*x)))
  cdf  (Normal u s _ _) x =
    0.5 * (1 + (erf ((x-u)/(s * (sqrt 2))) ))
  cdf (ChiSquared k _) x = (1/(gamma (kd/2))) * lig
    where
      kd = fromInteger $ toInteger k
      lig = lower_incomplete_gamma (kd /2) (x/2)
  cdf (Empirical b _) x = B.cdf b x
         

  cdf' (Uniform _) p = p
  cdf' (Exponential y _) p = -(log (1-p)) / y
  cdf' (Normal u s _ _) p =
    u + (s * (sqrt 2) * (inverf(2*p-1)))

  degreesOfFreedom (Uniform _) = 0
  degreesOfFreedom (Exponential _ _) = 1
  degreesOfFreedom (Normal _ _ _ _) = 2