{-# LANGUAGE ExistentialQuantification #-}

module Stochastic.Distributions.Continuous(
  mkExp
  ,mkNormal
  ,mkEmpirical
  ,mkChiSquared
  ,Dist(..)
  ,expTransform
  ,module Stochastic.Distribution.Continuous
  ,module Stochastic.Generators.Continuous
  ) where

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

import Stochastic.Generator
import Stochastic.Uniform
import Stochastic.Generators.Continuous

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 = Exponential Double
          | Normal Double Double (Maybe Double) 
          | ChiSquared Int 
          | Empirical B.Empirical 

data Sample = Sample Dist UniformRandom

instance RandomGen Sample where
  next g@(Sample (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@(Sample (Normal mean dev _) _) =
    let (x, g') = rand g in
    (truncate x, g')

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


mkEmpirical :: [Double] -> UniformRandom -> Sample
mkEmpirical samples = Sample $ Empirical (B.mkEmpirical samples) 

mkExp :: Double -> UniformRandom -> Sample
mkExp y = Sample $ Exponential y 

mkNormal :: Double -> Double -> UniformRandom -> Sample
mkNormal mean dev = Sample $ Normal mean dev Nothing 

mkChiSquared :: Int -> UniformRandom -> Sample
mkChiSquared d = Sample $ ChiSquared d

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


expTransform :: Double -> Double -> Double
expTransform y x = -(log $ x) / y

            
instance ContinuousSample Sample where
  entropy (Sample (Exponential y) u)     = entropy u
  entropy (Sample (Normal mean dev m) u) = (entropy u) + (fromMaybe 0 (fmap (\_->1) m))
  
  rand (Sample (Exponential y) u) =
    mapTuple (expTransform y) (Sample $ Exponential y) (rand u)
  rand (Sample (Normal mean dev m) uni) = f m
    where
      f (Just x) = (x, (Sample (Normal mean dev Nothing) uni'))
      f Nothing  = (y, (Sample (Normal mean dev (Just z)) uni'))
      ([u1, u2], uni') = rands 2 uni
      from_u g = mean + dev * (sqrt (-2 * (log u1))) * ( g (2 * pi * u2) )
      y = from_u (sin)
      z = from_u (cos)

instance ContinuousDistribution Sample where
  cdf  (Sample d _) = cdf  d
  cdf' (Sample d _) = cdf' d
  pdf  (Sample d _) = pdf  d
  degreesOfFreedom (Sample d _) = degreesOfFreedom d

  
instance ContinuousDistribution Dist where
  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' (Exponential y) p = -(log (1-p)) / y
  cdf' (Normal u s _) p =
    u + (s * (sqrt 2) * (inverf(2*p-1)))

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