{-# OPTIONS_GHC -fglasgow-exts #-}

module MonadRandom (
    MonadRandom
  , getRandom
  , getRandomR
  , evalRandomT
  , evalRand
  , fromList
  ) where

import System.Random
import Data.Ratio
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 RandomT m a = RandomT { unRT :: StateT StdGen m a }
  deriving (Functor, Monad, MonadTrans)

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) => MonadRandom (RandomT m) where
  getRandom = RandomT $ liftState random
  getRandomR (x,y) = RandomT $ liftState (randomR (x,y))

evalRandomT :: (Monad m) => RandomT m a -> StdGen -> m a
evalRandomT x g = evalStateT (unRT x) g

newtype Rand a = Rand { unRand :: RandomT Identity a }
  deriving (Functor, Monad, MonadRandom)

evalRand :: Rand a -> StdGen -> a
evalRand x g = runIdentity (evalRandomT (unRand x) g)

fromList :: [(a,Rational)] -> Rand 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 (\(_,q) (y,t) -> (y, t+q)) xs -- cumulative weight
                 p <- fmap toRational $ getRandomR (0.0 :: Double,s)
                 return . fst . head . dropWhile (\(_,q) -> q < p) $ cs