{- Helpers for random generation.  -}
module RandomHelp 
    ( module System.Random
    , module Control.Monad.State
    , module RandomHelp
    ) where
import Geometry

import System.Random
import Control.Monad.State
import Data.List

randomRanges :: (Random a,RandomGen g) => [a] -> State g a
randomRanges = join . takeOne . f
  where 
    f (x:y:ys) = state (randomR (x,y)) : f ys
    f  _ = []

takeOne :: RandomGen g => [a] -> State g a
takeOne xs = state (randomR (0,length xs - 1)) >>= (\i -> return (xs !! i))

takeOneWeighted :: (RandomGen g, Random b, Ord b, Num b) => [b] -> [a] -> State g a
takeOneWeighted ws xs = state (randomR (0, sum ws)) >>= (\w -> return (xs !! i w ws))
  where 
    i y (z:zs) 
        | y <= z = 0
        | otherwise = 1 + i (y-z) zs
    i _ _ = 0

takeOneMore :: RandomGen g => ([a],[a]) -> State g ([a],[a])
takeOneMore ( _,[]) = error "trying to takeOneMore from empty list"
takeOneMore (xs,ys) = do
    i <- state $ randomR (0,length ys - 1)
    let (zs, w:ws) = splitAt i ys
    return (w:xs, zs ++ ws)

takeNMore :: RandomGen g => Int -> ([a],[a]) -> State g ([a],[a])
takeNMore n p = foldl' (flip $ const (>>= takeOneMore)) (return p) [1..n]

takeN :: RandomGen g => Int -> [a] -> State g [a]
takeN 0  _ = return []
takeN i xs = fst <$> takeNMore i ([],xs)

-- | Randomly shuffle a list.
shuffle :: RandomGen g => [a] -> State g [a]
shuffle xs = do 
    rands <- forM [0..length xs-1] $ \i -> state $ randomR (0,i)
    let f ys rand = let (as,b:bs) = splitAt rand ys
                    in  (as ++ bs, b)
    let (_,zs) = mapAccumR f xs rands
    return zs

-- | Randomly shuffle the tail of a list, not safe.
shuffleTail :: RandomGen g => [a] -> State g [a]
shuffleTail (x:xs) = (x :) <$> shuffle xs
shuffleTail _ = undefined

-- select elements from a list randomly
-- each element has the same independent chance of being selected
randomSelectionFromList :: RandomGen g => Float -> [a] -> State g [a]
randomSelectionFromList = filterM . const . randProb

randProb :: RandomGen g => Float -> State g Bool
randProb p = do 
    p1 <- state $ randomR (0,1)
    return (p1 < p)

randInCirc :: RandomGen g => Float -> State g Point2
randInCirc = flip randInArc (2*pi)

randInArc :: RandomGen g => Float -> Float -> State g Point2
randInArc = randInArcStrip 0

randInArcStrip :: RandomGen g => Float -> Float -> Float -> State g Point2
randInArcStrip minrad maxRad maxangle = do 
    rad <- state $ randomR (minrad,maxRad)
    ang <- state $ randomR (0,maxangle)
    return $ rad *.* unitVectorAtAngle ang

randOnUnitSphere :: RandomGen g => State g Point3
randOnUnitSphere = do
    z <- state $ randomR (negate 1,1)
    longitude <- state $ randomR (0, 2*pi)
    let (V2 x y) = sqrt (1 - z^(2::Int)) *.* unitVectorAtAngle longitude
    return (V3 x y z)

randOnHemisphere :: RandomGen g => State g Point3
randOnHemisphere = do
    z <- state $ randomR (0,1)
    longitude <- state $ randomR (0, 2*pi)
    let (V2 x y) = sqrt (1 - z^(2::Int)) *.* unitVectorAtAngle longitude
    return (V3 x y z)

randInHemisphere :: RandomGen g => State g Point3
randInHemisphere = do
    p <- randOnHemisphere
    r <- state $ randomR (0,1)
    return $ r *.*.* p

randInRect :: RandomGen g => Float -> Float -> State g Point2
randInRect w h = do     
    x <- state $ randomR (0,w)
    y <- state $ randomR (0,h)
    return (V2 x y)

maybeTakeOne :: RandomGen g => [a] -> State g (Maybe a)
maybeTakeOne [] = return Nothing
maybeTakeOne xs = state (randomR (0,length xs - 1)) >>= (\i -> return (Just (xs !! i)))

randsSpread :: RandomGen g => (Float,Float) -> Int -> State g [Float]
randsSpread (a,b) i
    | i <= 0    = error "tried to take <= 0 randsSpread"
    | otherwise = zipWith (+) [a+x,a+2*x..] <$> replicateM i (state $ randomR (0,x))
  where
    x = (b-a)/fromIntegral i

randsOnCirc :: RandomGen g => Int -> State g [Float]
randsOnCirc = randsSpread (0,2*pi)
--randsOnCirc i
--    | i <= 0    = error "tried to take <= 0 randsOnCirc"
--    | otherwise = zipWith (+) [x,2*x..] <$> replicateM i (state $ randomR (0,x))
--  where
--    x = 2*pi/fromIntegral i