{-# LANGUAGE BangPatterns #-}
module Stochastic.Uniform(xorshift128plus,
                          UniformRandom,
                          nWayAllocate,
                          splitAllocate,
                          RandomGen(..)) where

import Stochastic.Generator
import Stochastic.Generators.Continuous as GC
import Stochastic.Generators.Discrete   as GD
import Stochastic.Distribution.Continuous
import Stochastic.Distribution.Discrete
import Data.Word
import Data.Bits
import Data.Typeable
import Control.Exception(throw, Exception)
import Data.Int
import GHC.Stack(errorWithStackTrace)
import System.Random(RandomGen(..))

data UniformRandom = XorShift128Plus Word64 Word64 Integer

data EntropyExhausted = EntropyExhausted
                      deriving(Eq, Typeable)

instance Exception EntropyExhausted where
instance Show EntropyExhausted where
  show e = "EntropyExhausted"

{- |
For information on the performance of the xorshift-128-plus PRNG, please see: <http://vigna.di.unimi.it/ftp/papers/xorshiftplus.pdf Vigna et al.>
-}
xorshift128plus :: Integer -> UniformRandom
xorshift128plus seed = XorShift128Plus high low entropy
  where
    high = fromIntegral seed
    low  = fromIntegral seed
    entropy = (2^127)

nWayAllocate :: (Integral a, Integral b) => a -> b -> UniformRandom -> ([UniformRandom], UniformRandom)
nWayAllocate _ 0 g0 = ([], g0)
nWayAllocate size n g0 = ((g1:gs), g3)
  where
    !(gs,g3) = nWayAllocate size (n-1) g2
    !(g1,g2) = splitAllocate size g0

splitAllocate :: Integral a => a -> UniformRandom -> (UniformRandom, UniformRandom)
splitAllocate count g@(XorShift128Plus high low entropy) =
    ((XorShift128Plus high low (toInteger count)), g')
    where
      !g' = step (next) (count) g


instance RandomGen UniformRandom where
  next (XorShift128Plus high low entropy) 
    | entropy == 0 = errorWithStackTrace "Entropy Exausted"
    | otherwise    = final 
    where
      -- eagerly evaluate this function, retain no intermediaries or we might blow the stack
      final    = (ret, XorShift128Plus high' low' entropy')
      !ret     = fromInteger $ toInteger $ high' + high
      x        = low `xor` (low `shift` 23)
      high'    = x `xor` high `xor` (x `shift` (-17)) `xor` (high `shift` (-26))
      low'     = high
      entropy' = entropy - 1
  split g@(XorShift128Plus high low entropy) =
    ((XorShift128Plus high low entropy'), g')
    where
      !entropy' = (2^32)
      !g' = step (next) (entropy') g

intRange = (toInteger (maxBound :: Int)) - (toInteger (minBound :: Int))

instance ContinuousDistribution UniformRandom where
  cdf  _ x   = x
  cdf' _ p   = p
  pdf  _ a b = b - a
  degreesOfFreedom _ = 0
  
instance DiscreteDistribution   UniformRandom where
  cdf _ i  = (fromIntegral i) / (fromIntegral intRange)
  cdf' _ p = (truncate (p * (fromIntegral intRange))) + (minBound :: Int)
  pmf _ i = 1.0 / (fromIntegral intRange)

intWordDouble :: Int -> Double
intWordDouble x =
  -- We use 53 bits of randomness corresponding to the 53 bit significand:
  ((fromIntegral (mask53 .&. (fromIntegral x::Int64)) :: Double)  
	   /  fromIntegral twoto53)
  where 
    twoto53 = (2::Int64) ^ (53::Int64)
    mask53 = twoto53 - 1

    
instance ContinuousSample UniformRandom where
  rand    g = let (x, g') = next g in (intWordDouble x, g')
  entropy (XorShift128Plus high low e) = fromIntegral $ min e (2^63)
  
instance DiscreteSample   UniformRandom where
  rand    = next
  entropy (XorShift128Plus high low e) = fromIntegral $ min e (2^63)
  

step :: Integral b => (g -> (a,g)) -> b -> g -> g
step f 0 g'' = g''
step f n g'' = step (f) (n-1) $! (snd $ f g'')