{-# LANGUAGE BangPatterns         #-}
{-# LANGUAGE DeriveDataTypeable   #-}
{-# LANGUAGE DeriveFoldable       #-}
{-# LANGUAGE DeriveFunctor        #-}
{-# LANGUAGE DeriveGeneric        #-}
{-# LANGUAGE DeriveTraversable    #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE InstanceSigs         #-}
{-# LANGUAGE ScopedTypeVariables  #-}
{-# LANGUAGE TupleSections        #-}
{-# LANGUAGE TypeApplications     #-}
{-# LANGUAGE UndecidableInstances #-}

-- |
-- Module      : System.Random.OneLiner
-- Description : Derived methods for Random.
-- Copyright   : (c) Justin Le 2018
-- License     : BSD-3
-- Maintainer  : justin@jle.im
-- Stability   : unstable
-- Portability : portable
--
-- Derived methods for 'Random', using "Generics.OneLiner" and
-- "GHC.Generics".
--
-- Can be used for any types (deriving 'Generic') made with a single
-- constructor, where every field is an instance of 'Random'.
--
-- Also includes a newtype wrapper that imbues any such data type with
-- instant 'Random' instances, which can one day be used with /DerivingVia/
-- syntax to derive instances automatically.
--

module System.Random.OneLiner (
  -- * Single constructor
  -- ** Newtype wrapper
    GRandom(..)
  -- ** Generics-derived methods
  , gRandomR
  , gRandom
  , gRandomRs
  , gRandoms
  , gRandomRIO
  , gRandomIO
  -- * Multiple constructor
  -- ** Newtype wrapper
  , GRandomSum(..)
  -- ** Generics-derived methods
  , gRandomRSum
  , gRandomSum
  , gRandomRSums
  , gRandomSums
  , gRandomRIOSum
  , gRandomIOSum
  ) where

import           Control.Monad
import           Data.Coerce
import           Data.Data
import           Data.Functor.Compose
import           Data.List.NonEmpty   (NonEmpty(..))
import           Data.Maybe
import           GHC.Exts             (build)
import           GHC.Generics
import           Generics.OneLiner
import           System.Random
import qualified Data.List.NonEmpty   as NE

-- | If @a@ is a data type with a single constructor whose fields are all
-- instances of 'Random', then @'GRandom' a@ has a 'Random' instance.
--
-- Will one day be able to be used with /DerivingVia/ syntax, to derive
-- instances automatically.
--
-- Only works with data types with single constructors.  If you need it to
-- work with types of multiple constructors, consider 'GRandomSum'.
newtype GRandom a = GRandom { getGRandom :: a }
  deriving (Eq, Ord, Show, Read, Data, Generic, Functor, Foldable, Traversable)

instance ( ADTRecord a
         , Constraints a Random
         )
      => Random (GRandom a) where
    randomR :: forall g. RandomGen g => (GRandom a, GRandom a) -> g -> (GRandom a, g)
    randomR = coerce (gRandomR @a @g)
    {-# INLINE randomR #-}
    random :: forall g. RandomGen g => g -> (GRandom a, g)
    random = coerce (gRandom @a @g)
    {-# INLINE random #-}
    randomRs :: forall g. RandomGen g => (GRandom a, GRandom a) -> g -> [GRandom a]
    randomRs = coerce (gRandomRs @a @g)
    {-# INLINE randomRs #-}
    randoms :: forall g. RandomGen g => g -> [GRandom a]
    randoms = coerce (gRandoms @a @g)
    {-# INLINE randoms #-}
    randomRIO :: (GRandom a, GRandom a) -> IO (GRandom a)
    randomRIO = coerce (gRandomRIO @a)
    {-# INLINE randomRIO #-}
    randomIO :: IO (GRandom a)
    randomIO = coerce (gRandomIO @a)
    {-# INLINE randomIO #-}

-- | 'randomR' implemented by sequencing 'randomR' between all components
--
-- Requires the type to have only a single constructor.
gRandomR
    :: forall a g. (ADTRecord a, Constraints a Random, RandomGen g)
    => (a, a) -> g -> (a, g)
gRandomR (l, u) = runState $
    dialgebra @Random
      (State . randomR . dePair)
      (Pair l u)
{-# INLINE gRandomR #-}

-- | 'random' implemented by sequencing 'random' for all components.
--
-- Requires the type to have only a single constructor.
gRandom
    :: forall a g. (ADTRecord a, Constraints a Random, RandomGen g)
    => g -> (a, g)
gRandom = runState $ createA' @Random (State random)
{-# INLINE gRandom #-}

-- | 'randomRs' implemented by repeatedly calling 'gRandomR'.
gRandomRs
    :: forall a g. (ADTRecord a, Constraints a Random, RandomGen g)
    => (a, a) -> g -> [a]
gRandomRs ival g = build (\cons _nil -> buildRandoms cons (gRandomR ival) g)
{-# INLINE gRandomRs #-}

-- | 'randoms' implemented by repeatedly calling 'gRandom'.
gRandoms
    :: forall a g. (ADTRecord a, Constraints a Random, RandomGen g)
    => g -> [a]
gRandoms g = build (\cons _nil -> buildRandoms cons gRandom g)
{-# INLINE gRandoms #-}

-- | 'randomRIO' implemented by calling 'gRandomR' on the global seed.
gRandomRIO
    :: forall a. (ADTRecord a, Constraints a Random)
    => (a, a) -> IO a
gRandomRIO range = getStdRandom (gRandomR range)
{-# INLINE gRandomRIO #-}

-- | 'randomIO' implemented by calling 'gRandom' on the global seed.
gRandomIO
    :: forall a. (ADTRecord a, Constraints a Random)
    => IO a
gRandomIO = getStdRandom gRandom
{-# INLINE gRandomIO #-}

-- | If @a@ is a data type whose fields are all instances of 'Random', then
-- @'GRandom' a@ has a 'Random' instance.
--
-- Will one day be able to be used with /DerivingVia/ syntax, to derive
-- instances automatically.
--
-- A version of 'GRandom' that works for data types with multiple
-- constructors.  If your type has only one constructor, it might be more
-- performant to use 'GRandom'.
--
-- Note that the "ranged" variants are partial: if given a range of items
-- made with different constructors, will be 'error'!
newtype GRandomSum a = GRandomSum { getGRandomSum :: a }
  deriving (Eq, Ord, Show, Read, Data, Generic, Functor, Foldable, Traversable)

instance ( ADT a
         , Constraints a Random
         )
      => Random (GRandomSum a) where
    randomR :: forall g. RandomGen g => (GRandomSum a, GRandomSum a) -> g -> (GRandomSum a, g)
    randomR = coerce (gRandomRSum @a @g)
    {-# INLINE randomR #-}
    random :: forall g. RandomGen g => g -> (GRandomSum a, g)
    random = coerce (gRandomSum @a @g)
    {-# INLINE random #-}
    randomRs :: forall g. RandomGen g => (GRandomSum a, GRandomSum a) -> g -> [GRandomSum a]
    randomRs = coerce (gRandomRSums @a @g)
    {-# INLINE randomRs #-}
    randoms :: forall g. RandomGen g => g -> [GRandomSum a]
    randoms = coerce (gRandomSums @a @g)
    {-# INLINE randoms #-}
    randomRIO :: (GRandomSum a, GRandomSum a) -> IO (GRandomSum a)
    randomRIO = coerce (gRandomRIOSum @a)
    {-# INLINE randomRIO #-}
    randomIO :: IO (GRandomSum a)
    randomIO = coerce (gRandomIOSum @a)
    {-# INLINE randomIO #-}

-- | 'randomR' implemented by sequencing 'randomR' between all components.
--
-- If given a range of items made with different constructors, will be
-- 'error'!
gRandomRSum
    :: forall a g. (ADT a, Constraints a Random, RandomGen g)
    => (a, a) -> g -> (a, g)
gRandomRSum (l, u) = runState . fromMaybe (error badbad) . getCompose $
    zipWithA @Random (\l' u' -> Compose (Just (State (randomR (l', u')))))
      l u
  where
    badbad = "gRandomRSum: Constructors in range do not match."
{-# INLINE gRandomRSum #-}

-- | 'random' implemented by selecting a random constructor and sequencing
-- 'random' for all components.
gRandomSum
    :: forall a g. (ADT a, Constraints a Random, RandomGen g)
    => g -> (a, g)
gRandomSum = case options of
    Nothing   -> (error "gRandomSum: Uninhabited type",)
    Just opts -> runState (join (reservoir opts))
  where
    options = NE.nonEmpty . getCompose $ createA @Random @a $
        Compose [State random]
{-# INLINE gRandomSum #-}

-- | 'randomRs' implemented by repeatedly calling 'gRandomRSum'.
--
-- If given a range of items made with different constructors, will be
-- 'error'!
gRandomRSums
    :: forall a g. (ADT a, Constraints a Random, RandomGen g)
    => (a, a) -> g -> [a]
gRandomRSums ival g = build (\cons _nil -> buildRandoms cons (gRandomRSum ival) g)
{-# INLINE gRandomRSums #-}

-- | 'randoms' implemented by repeatedly calling 'gRandomSum'.
gRandomSums
    :: forall a g. (ADT a, Constraints a Random, RandomGen g)
    => g -> [a]
gRandomSums g = build (\cons _nil -> buildRandoms cons gRandomSum g)
{-# INLINE gRandomSums #-}

-- | 'randomRIO' implemented by calling 'gRandomRSum' on the global seed.
--
-- If given a range of items made with different constructors, will be
-- 'error'!
gRandomRIOSum
    :: forall a. (ADT a, Constraints a Random)
    => (a, a) -> IO a
gRandomRIOSum range = getStdRandom (gRandomRSum range)
{-# INLINE gRandomRIOSum #-}

-- | 'randomIO' implemented by calling 'gRandom' on the global seed.
gRandomIOSum
    :: forall a. (ADT a, Constraints a Random)
    => IO a
gRandomIOSum = getStdRandom gRandomSum
{-# INLINE gRandomIOSum #-}


data Pair a = Pair !a !a
    deriving Functor

dePair :: Pair a -> (a, a)
dePair (Pair x y) = (x, y)
{-# INLINE dePair #-}

newtype State s a = State { runState :: s -> (a, s) }
    deriving Functor

instance Applicative (State s) where
    pure x = State (x,)
    {-# INLINE pure #-}
    sf <*> sx = State $ \s0 ->
        let (f, !s1) = runState sf s0
            (x, !s2) = runState sx s1
        in  (f x, s2)
    {-# INLINE (<*>) #-}

instance Monad (State s) where
    return x = State (x,)
    {-# INLINE return #-}
    sx >>= f = State $ \s0 ->
      let (x, !s1) = runState sx s0
      in  runState (f x) s1
    {-# INLINE (>>=) #-}

buildRandoms :: RandomGen g
             => (a -> as -> as)  -- ^ E.g. '(:)' but subject to fusion
             -> (g -> (a,g))     -- ^ E.g. 'random'
             -> g                -- ^ A 'RandomGen' instance
             -> as
buildRandoms cons rand = go
  where
    -- The seq fixes part of #4218 and also makes fused Core simpler.
    go g = x `seq` (x `cons` go g') where (x,g') = rand g
{-# INLINE buildRandoms #-}

-- | Select a random item from a non-empty list in constant memory, using
-- only a single traversal, using reservoir sampling.
reservoir :: RandomGen g => NE.NonEmpty a -> State g a
reservoir (x :| xs) = go 2 x xs
  where
    go _  y []     = pure y
    go !i y (z:zs) = do
      j <- State $ randomR @Int (1, i)
      if j <= 1
        then go (i + 1) z zs
        else go (i + 1) y zs
{-# INLINE reservoir #-}