module Stochastic.Generator where
import Control.Concurrent.MVar
import Control.Monad.State.Lazy

type Gen g a = (g -> (a,g))
data IOGen g a = IOGen (g -> (a, g)) (MVar g)

liftGen :: (g -> (a, g)) -> g -> IO (IOGen g a)
liftGen f g = do
  var <- newMVar g
  return $ IOGen f var

nextIO :: IOGen g a -> IO a
nextIO (IOGen f var) = do
  val <- takeMVar var
  let (x, g') = f val
  putMVar var g'
  return x

foldGenWhile :: (g -> (a,g))
          -> (b -> a -> b)
          -> b
          -> (b -> Bool)
          -> (g -> ([a], g))
foldGenWhile nxt f zz p = h zz
  where
    h z g0 
      | not (p z) = ([], g0)
      | otherwise = (x:xs, g2)
      where
        (xs, g2) = h (f z x) g1
        (x, g1)  = nxt g0

genWhile :: (g -> (a, g)) -> (a -> Bool) -> (g -> ([a], g))
genWhile nxt p = h
  where
    h g0 = let (x, g1) = nxt g0 in
            let (xs, g2) = h g1 in
            if (p x) then (x:xs, g2) else ([], g1)

genTake :: (Eq b, Num b) => (g -> (a,g)) -> b -> (g -> ([a], g))
genTake f 0 g0 = ([], g0)
genTake f n g0 = ((x:xs), g2)
  where
    (x, g1) = f g0
    (xs, g2) = genTake f (n-1) g1

dropGen :: (Eq b, Num b) => (g -> (a,g)) -> b -> g -> g
dropGen f = d
  where
    d 0 g0 = g0
    d n g0 = d (n-1) $! (snd $ f g0)



dropIO :: IOGen g a -> Integer -> IO ()
dropIO _ 0 = return ()
dropIO ioG n = do
  nextIO ioG
  dropIO ioG (n-1)