-- | This module contains a permuted linear congruential pseudorandom number
-- generator, as described by M.E. O'Neill (<pcg-random.org>). This version
-- holds two Word64 values and outputs a Word32 of randomness each time you use
-- it. Compared to the StdGen type from System.Random, it's around a 2.5x to 3x
-- speedup on a 64-bit system. Runs somewhat slower than StdGen on a 32-bit
-- system (the Word64 values must be emulated after all), but there aren't too
-- many of those in play these days anyway.
--
-- The first Word64 is the "state", which changes each step of the generator.
-- The second Word64 is the "inc", which stays fixed from step to step. The inc
-- must always be odd for the gerator to work properly (so there are 2^63
-- possible inc values), though the smart constructors will ensure this. The
-- state value eventually returns to it's initial position after 2^64 uses, and
-- each inc value is a different sequence of states.
module Data.PCGen (
    PCGen,
    mkPCGen,
    mkPCGenDetailed,
    stepGen
    ) where

-- base
import Data.Bits
import Data.Word (Word32,Word64)
import Data.Int (Int32)
-- random
import System.Random

-- | The PCGen data type. You generally create values of this type with mkPCGen
-- and use them with stepGen or the next method of the RandomGen instance.
-- Possibly in combination with a State, StateT, MonadRandom, etc.
--
-- Note that (at the moment) the Read instance is simply derived by the
-- compiler, so it won't be able to ensure an odd inc value. It's guaranteed
-- that any String you get from the Show instance of this data type (also
-- derived) will be correct as long as the original value was made with
-- 'mkPCGen', but otherwise you must be careful. This will be fixed later.
data PCGen = PCGen !Word64 !Word64
    -- Note: GHC unpacks small fields by default when they're strict, and if the
    -- user for some reason *did* that features off, then we should respect
    -- their wishes and not unpack our crap in their space. So we deliberately
    -- leave out the UNPACK directive on these fields.
    deriving (Read, Show, Eq, Ord)
-- TODO: Fix the above thing where Read won't ensure an odd inc value by writing
-- a custom Read instance.

-- | Creates a new PCGen value by using the Integral given as both the "state"
-- and "inc" values for the generator. If the value given isn't odd, then it's
-- bumped to the next higher odd value for the inc. The state of the generator
-- is then advanced once, because otherwise the first result tends to be 0 with
-- human picked seeds. See also 'mkPCGenDetailed'.
mkPCGen :: Integral i => i -> PCGen
mkPCGen n = let
    n' = fromIntegral n
    in snd $ stepGen $ PCGen n' (n' .|. 1)

-- | Creates a PCGen using the specified state and inc values and returns it
-- without an initial generator use. It still bumps the inc value up to the next
-- odd value if an even value is given.
mkPCGenDetailed :: Word64 -> Word64 -> PCGen
mkPCGenDetailed st inc = PCGen st (inc .|. 1)

-- | Advances the given generator one step, giving back a Word32 of output and
-- the resultant generator as well. This is the most basic way to advance the
-- generator, and if this is all you're going for then you might want to look at
-- the RandomGen instance and its 'next' method. Other functions in this module
-- let you generate values in batches and so forth, such as for rolling dice.
stepGen :: PCGen -> (Word32, PCGen)
stepGen (PCGen state inc) = let
    xorshifted = ((state `shiftR` 18) `xor` state) `shiftR` 27
    rot = fromIntegral $ state `shiftR` 59
    out = fromIntegral $ (xorshifted `shiftR` rot) .|.
        (xorshifted `shiftL` ((-rot) .&. 31))
    newState = state * 6364136223846793005 + inc
    in (out, PCGen newState inc)

instance RandomGen PCGen where
    -- 'next' specifies that the output is an Int value, so we convert the
    -- Word32 bits into the Int32 range before converting that into an Int and
    -- returning. If we convert the Word32 directly to Int and Int is Int64,
    -- then our outputs won't match up properly with what we declare with the
    -- genRange function. Of course, all this number conversion stuff between
    -- Integral types is pretty much noop nonsense once we compile, but Haskell
    -- likes to be very precise with types after all.
    next gen = let
        (outWord, nextGen) = stepGen gen
        outInt = fromIntegral (fromIntegral outWord :: Int32) :: Int
        in (outInt, nextGen)

    -- Similar to the above, the range of a PCGen is 32-bits of output per step,
    -- and so we use the bounds of Int32, but we have to convert that into Int
    -- to conform with the spec of the typeclass.
    genRange _ = (fromIntegral (minBound :: Int32),
                    fromIntegral (maxBound :: Int32))

    -- The only real spec here is that the two result generators be dissimilar
    -- from each other and also from the input generator. So we just do some
    -- nonsense shuffling around to achieve that.
    split gen@(PCGen state inc) = let -- no statistical foundation for this!
        (q,nGen1@(PCGen sa ia)) = stepGen gen
        (w,nGen2@(PCGen sb ib)) = stepGen nGen1
        (e,nGen3@(PCGen sc ic)) = stepGen nGen2
        (r,nGen4@(PCGen sd id)) = stepGen nGen3
        stateA = sd `rotateR` 5
        stateB = sd `rotateR` 3
        incA = ((fromIntegral q) `shiftL` 32) .|. (fromIntegral w)
        incB = ((fromIntegral e) `shiftL` 32) .|. (fromIntegral r)
        outA = PCGen stateA (incA .|. 1)
        outB = PCGen stateB (incB .|. 1)
        in (outA, outB)
        -- TODO: This could probably be faster while still conforming to spec.

instance Random PCGen where
    -- produces a random PCGen value.
    random gen = let
        (x,newGen) = random gen
        in (mkPCGen (x::Word),newGen)
    
    -- produces a random PCGen value but ignores the input since it would
    -- nonsensical to try and use them as the "bounds" of anything.
    randomR (_,_) gen = let
        (x,newGen) = random gen
        in (mkPCGen (x::Word),newGen)