{-# OPTIONS_GHC -fglasgow-exts #-}
 
module MonadRandom (
    MonadRandom,
    getRandom,
    getRandomR,
    evalRandomT,
    evalRand,
    evalRandIO,
    fromList,
    Rand, RandomT -- but not the data constructors
    ) where
 
import System.Random
import Control.Monad.State
import Control.Monad.Identity
 
class (Monad m) => MonadRandom m where
    getRandom :: (Random a) => m a
    getRandomR :: (Random a) => (a,a) -> m a
 
newtype (RandomGen g) => RandomT g m a = RandomT { unRT :: StateT g m a }
    deriving (Functor, Monad, MonadTrans, MonadIO)
 
liftState :: (MonadState s m) => (s -> (a,s)) -> m a
liftState t = do v <- get
                 let (x, v') = t v
                 put v'
                 return x
 
instance (Monad m, RandomGen g) => MonadRandom (RandomT g m) where
    getRandom = (RandomT . liftState) random
    getRandomR (x,y) = (RandomT . liftState) (randomR (x,y))
 
evalRandomT :: (Monad m, RandomGen g) => RandomT g m a -> g -> m a
evalRandomT x g = evalStateT (unRT x) g
 
-- Boring random monad :)
newtype Rand g a = Rand { unRand :: RandomT g Identity a }
    deriving (Functor, Monad, MonadRandom)
 
evalRand :: (RandomGen g) => Rand g a -> g -> a
evalRand x g = runIdentity (evalRandomT (unRand x) g)
 
evalRandIO :: Rand StdGen a -> IO a
evalRandIO x = getStdRandom (runIdentity . runStateT (unRT (unRand x)))
 
fromList :: (MonadRandom m) => [(a,Rational)] -> m a
fromList [] = error "MonadRandom.fromList called with empty list"
fromList [(x,_)] = return x
fromList xs = do let s = fromRational $ sum (map snd xs) -- total weight
                     cs = scanl1 (\(x,q) (y,s) -> (y, s+q)) xs -- cumulative weight
                 p <- liftM toRational $ getRandomR (0.0,s)
                 return $ fst $ head $ dropWhile (\(x,q) -> q < p) cs

instance (MonadRandom m) => MonadRandom (StateT s m) where
    getRandom = lift getRandom
    getRandomR r = lift $ getRandomR r
 
instance (MonadRandom m, Monoid w) => MonadRandom (WriterT w m) where
    getRandom = lift getRandom
    getRandomR r = lift $ getRandomR r
 
instance (MonadRandom m) => MonadRandom (ReaderT r m) where
    getRandom = lift getRandom
    getRandomR r = lift $ getRandomR r
 
 
instance (MonadState s m, RandomGen g) => MonadState s (RandomT g m) where
    get = lift get
    put s = lift $ put s
 
instance (MonadReader r m, RandomGen g) => MonadReader r (RandomT g m) where
    ask = lift ask
    local f m = RandomT $ local f (unRT m)
 
instance (MonadWriter w m, RandomGen g, Monoid w) => MonadWriter w (RandomT g m) where
    tell w = lift $ tell w
    listen m = RandomT $ listen (unRT m)
    pass m = RandomT $ pass (unRT m)