module RandT (
  RandT,
  evalRandT,
  generate,
  rnd, perhaps, upto, choice,
  random_field
) where

import Data.Array
import Data.Ratio
import Data.Word
import Control.Monad.State
import System.Random

import Util

newtype RandT m a = RandT (StateT [Int] m a)

-- Syntax simplifier
{-# INLINE derand #-}
derand :: RandT m a -> StateT [Int] m a
derand (RandT m) = m

-- Define instances that can be delegated directly to StateT.

instance (Monad m) => Functor (RandT m) where
  fmap f (RandT m) = RandT $ fmap f m

instance (Monad m) => Monad (RandT m) where
  return a = RandT $ return a
  (RandT m) >>= k = RandT $ m >>= (derand . k)
  fail str = RandT $ fail str

instance (MonadPlus m) => MonadPlus (RandT m) where
  mzero = RandT mzero
  (RandT m) `mplus` (RandT n) = RandT $ m `mplus` n

instance (MonadFix m) => MonadFix (RandT m) where
  mfix f = RandT $ mfix (derand . f)

instance MonadTrans RandT where
  lift m = RandT $ lift m

instance (MonadIO m) => MonadIO (RandT m) where
  liftIO = lift . liftIO

evalRandT :: (Monad m, RandomGen g) => RandT m a -> g -> m a
evalRandT (RandT m) g = evalStateT m (randoms g)

-- Switch child monads while retaining the same RandT context.
generate :: (Monad m) => RandT (State a) () -> RandT m a
generate (RandT n) = do
  rs <- RandT $ get
  -- The result of the execStateT is a (State a) [Int] monad.
  let (rs', gen) = runState (execStateT n rs) undefined
  RandT $ put rs'
  return gen

-- Take the first number from the stream of random numbers
pop :: (Monad m) => RandT m Int
pop = do
  top <- RandT $ gets head
  RandT $ modify tail
  return top

-- Return a random number in the given range (endpoints included).
-- Requires that high >= low.
rnd :: (Monad m) => (Int, Int) -> RandT m Int
rnd (low, high) = do
  r <- wrnd (fromIntegral (high - low + 1))
  return (low + fromIntegral r)

-- Return a randum number in the range (0, lim-1).  The obvious approach
-- of using r `mod` lim would bias the result slightly downward if
-- lim does not divide evenly into maxBound, so we sometimes consume
-- more than one token in order to get it right.
wrnd :: (Monad m) => Word -> RandT m Word
wrnd lim = do
  r <- pop
  let m = (fromIntegral r) `mod` lim
      base = (fromIntegral r) - m
      maxm = lim - 1
  -- check if base + maxm would have gone over maxBound.  If so,
  -- we need to try again with a new token, because this one didn't
  -- cover the full range of m.
  if maxBound - base < maxm
    then wrnd lim
    else return m

perhaps :: (Monad m) => Ratio Int -> RandT m () -> RandT m ()
perhaps ratio m = do
  r <- rnd (1, denominator ratio)
  if r <= numerator ratio
    then m
    else return ()

upto :: (Monad m) => Int -> RandT m () -> RandT m ()
upto maxtimes m = do
  times <- rnd (0, maxtimes)
  sequence_ (replicate times m)

choice :: (Monad m) => [RandT m a] -> RandT m a
choice ms = do
  c <- rnd (0, length ms - 1)
  ms !! c

random_field :: (Monad m, Ix i) => (Int, Int) -> (i, i) -> RandT m (i -> Int)
random_field (low, high) (from, to) = do
  rs <- sequence $ replicate (rangeSize (from, to)) $ rnd (low, high)
  let la = listArray (from, to) rs
  return (lookupA la 0)