{-# 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 Control.Comonad.Random
    ( module Control.Comonad
    , module System.Random

      -- * Example
      -- $Example

      -- * Data Type
    , Rand ()

      -- ** Creation
    , mkRand
    , mkRandR

      -- ** Transformations
    , next
    , left
    , right

      -- ** Convenience
    , extracts
    )
    where

import Control.Applicative
import Control.Comonad
import Control.Comonad.Cofree
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 }
    deriving (Functor, Copointed, Comonad)

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.

-}