```{-# LANGUAGE GeneralizedNewtypeDeriving, Rank2Types #-}

-- | This module provides a comonadic interface to random values. In
--   some situations, this may be more natural than a monadic
--   approach.
, module System.Random

-- * Example
-- \$Example

-- * Data Type
, Rand ()

-- ** Creation
, mkRand
, mkRandR

-- ** Transformations
, next
, left
, right

-- ** Convenience
, extracts
)
where

import Control.Applicative
import Control.Functor.Pointed
import Data.Function
import System.Random hiding (next)

newtype Three a = Three (a, a, a)

fst3 :: Three a -> a
fst3 (Three (x, _, _)) = x

snd3 :: Three a -> a
snd3 (Three (_, y, _)) = y

thd3 :: Three a -> a
thd3 (Three (_, _, z)) = z

instance Functor Three where
fmap f (Three ~(x, y, z)) = Three (f x, f y, f z)

instance Applicative Three where
pure x = Three (x, x, x)
Three ~(f, g, h) <*> Three ~(x, y, z) = Three (f x, g y, h z)

-- | A memoized supply of values
newtype Rand a = Rand { unRand :: Cofree Three a }

inRand :: (Cofree Three a -> Cofree Three b) -> (Rand a -> Rand b)
inRand = (Rand .) . (. unRand)

instance Applicative Rand where
pure x = let tree = cofree x \$ Three (tree, tree, tree)
in Rand tree
Rand a <*> Rand b = let ~(f, fs) = runCofree a
~(x, xs) = runCofree b
in Rand . cofree (f x) . fmap unRand .
liftA2 (<*>) (fmap Rand fs) . fmap Rand \$ xs

mkRandWith :: RandomGen g => (g -> (a, g)) -> g -> Rand a
mkRandWith f g = let ~(x, g') = f g
~(l, r) = split g
in Rand . cofree x . fmap (unRand . mkRandWith f) . Three \$ (l, g', r)

-- | Create a comonadic generator from a 'RandomGen'.
mkRand :: (RandomGen g, Random a) => g -> Rand a
mkRand = mkRandWith random

-- | Create a comonadic generator from a 'RandomGen' where the values
--   are limited to a given range.
mkRandR :: (RandomGen g, Random a) => (a, a) -> g -> Rand a
mkRandR = mkRandWith . randomR

inner :: (forall b . Three b -> b) -> (Rand a -> Rand a)
inner f = inRand \$ f . snd . runCofree

-- | Get the generator for the next value.
next :: Rand a -> Rand a
next = inner snd3

-- | Split the generator, returning the new left one.
left :: Rand a -> Rand a
left = inner fst3

-- | Split the generator, returning the new right one.
right :: Rand a -> Rand a
right = inner thd3

-- | Generate an infinite list of values by applying a function
--   repeatedly.
extracts :: Copointed f => (f a -> f a) -> f a -> [a]
extracts f = liftA2 (:) extract (extracts f . f)

{- \$Example

The following function generates an infinite list of dice throw
sums with @n@ dice.

> rolls :: RandomGen g => Int -> g -> [Int]
> rolls n =
>     extracts left          .  -- Extract an infinite list of the sums.
>     fmap (sum . take n)    .  -- Sum the first n values of each list.
>     extend (extracts next) .  -- Group them into lists of die values.
>     mkRandR (1,6)             -- Generate random die values.

One potential gotcha with this library is that a top-level @Rand@
that is extracted deeply could result in a space leak due to the
memoization. It's a good idea to try not to hold on to @Rand@s for
longer than necessary.

-}
```