module Random (setup, next, randomList) where

import System.Random hiding (next)
import Data.Array.IArray


-- Given a list of items and a list of their propabilities generate a tripel
-- consisting of the values vector, the alias vector and the relative propabilities
-- vector which is used in applications of next, etc.
setup :: (Ord a, IArray a2 a, IArray a1 e, Num a) => [e] -> [a] -> (a1 Int e, a1 Int e, a2 Int a)
setup xs ps = 
    let n  = length ps
        xv = listArray (0, n - 1) xs
        rv = listArray (0, n - 1) [fromIntegral n * p | p <- ps]
        (low, high) = splitAt 1 rv (indices rv) [] []
        (a, r) = calcAlias xv xv rv high low
    in
        (xv, a, r)
    where
        -- Return a pair of lists, the first consisting of elements lower than
        -- given threshold and the second with elements greater than threshold,
        -- equal elements are ignored.
        splitAt t v []     l h = (l, h)
        splitAt t v (i:is) l h = case v!i of
            x | x < t     -> splitAt t v is (i:l) h
              | x > t     -> splitAt t v is l (i:h)
              | otherwise -> splitAt t v is l h


-- Given an list of highs and a list of lows, calculate the alias vector and the relative
-- propabilities vector.
calcAlias :: (Ord e, Num e, IArray a e, Ix i, IArray a2 e1, IArray a1 e1) => a2 i e1 -> a1 i e1 -> a i e -> [i] -> [i] -> (a1 i e1, a i e)
calcAlias xv av rv []        _      = (av, rv)
calcAlias xv av rv _         []     = (av, rv)
calcAlias xv av rv hi@(h:hs) (l:ls) =
    let av' = av//[(l, xv!h)]
        rv' = rv//[(h, rv!h + rv!l - 1)]
    in
        if rv'!h >= 1
            then calcAlias xv av' rv' hi ls
            else calcAlias xv av' rv' hs (h:ls)
        

-- Generate a random item according to the given propability distribution as specified
-- by the given tripel which is the result of applying "setup" to a list of items and
-- a list of propabilities.
next :: (IArray a2 e1, IArray a e1, Ord e, IArray a1 e, RandomGen t, Random e) => (a Int e1, a2 Int e1, a1 Int e) -> t -> (e1, t)
next (xs, as, rs) g =
    let n = length $ indices xs
        (x1, g1) = randomR (0, n - 1) g
        (x2, g2) = random g1
        r = rs!x1
    in
        if x2 <= r 
            then (xs!x1, g2) 
            else (as!x1, g2)


-- Generate a infinite list of random items according to the specified propability
-- distribution as given by the triple that results from applying setup to a pair
-- of a list of items and a list of propabilities (see "setup" for details).
randomList :: (Random e, RandomGen t1, IArray a2 e, Ord e, IArray a t, IArray a1 t) => (a Int t, a1 Int t, a2 Int e) -> t1 -> [t]
randomList t g = 
    let (n, g') = next t g 
    in 
        n:randomList t g'