{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-|
Module      : IT.Random
Description : random generation of IT components
Copyright   : (c) Fabricio Olivetti de Franca, 2020
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com
Stability   : experimental
Portability : POSIX

-}
module IT.Random where

import IT

import System.Random
import Control.Monad.State
import qualified Data.IntMap.Strict as M

-- * Random expressions generation

-- | The type used for generating random values of 'a'
type Rnd a = State StdGen a

-- | Creates a random interaction with up to n variables. 
--   Guaranteed to choose at least one variable.
sampleInterMax :: Int               -- ^ problem dimension
               -> Int               -- ^ maximum number of variables in this interaction
               -> Int               -- ^ minimum exponent            
               -> Int               -- ^ maximum exponent
               -> Rnd Interaction   -- ^ Random interaction generator
sampleInterMax :: Int -> Int -> Int -> Int -> Rnd Interaction
sampleInterMax Int
dim Int
budget Int
minExp Int
maxExp = do [(Int, Int)]
es <- Int -> Int -> Int -> Int -> Rnd [(Int, Int)]
sampleInterMax' Int
dim Int
budget Int
minExp Int
maxExp
                                             Interaction -> Rnd Interaction
forall (m :: * -> *) a. Monad m => a -> m a
return (Interaction -> Rnd Interaction) -> Interaction -> Rnd Interaction
forall a b. (a -> b) -> a -> b
$ [(Int, Int)] -> Interaction
forall a. [(Int, a)] -> IntMap a
M.fromList [(Int, Int)]
es

sampleInterMax' :: Int -> Int -> Int -> Int -> Rnd [(Int, Int)]
sampleInterMax' :: Int -> Int -> Int -> Int -> Rnd [(Int, Int)]
sampleInterMax' Int
0   Int
_      Int
_      Int
_      = [(Int, Int)] -> Rnd [(Int, Int)]
forall (m :: * -> *) a. Monad m => a -> m a
return []
sampleInterMax' Int
_   Int
0      Int
_      Int
_      = [(Int, Int)] -> Rnd [(Int, Int)]
forall (m :: * -> *) a. Monad m => a -> m a
return []
sampleInterMax' Int
1   Int
_      Int
minExp Int
maxExp = do Int
e <- Int -> Int -> Rnd Int
sampleNZRng Int
minExp Int
maxExp
                                              [(Int, Int)] -> Rnd [(Int, Int)]
forall (m :: * -> *) a. Monad m => a -> m a
return [(Int
0,Int
e)]
sampleInterMax' Int
dim Int
budget Int
minExp Int
maxExp = do Bool
b <- Rnd Bool
toss
                                              if Bool
b
                                              then do Int
e  <- Int -> Int -> Rnd Int
sampleNZRng Int
minExp Int
maxExp
                                                      [(Int, Int)]
es <- Int -> Int -> Int -> Int -> Rnd [(Int, Int)]
sampleInterMax' (Int
dimInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) (Int
budgetInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
minExp Int
maxExp
                                                      [(Int, Int)] -> Rnd [(Int, Int)]
forall (m :: * -> *) a. Monad m => a -> m a
return ((Int
dimInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1,Int
e)(Int, Int) -> [(Int, Int)] -> [(Int, Int)]
forall a. a -> [a] -> [a]
:[(Int, Int)]
es)
                                              else Int -> Int -> Int -> Int -> Rnd [(Int, Int)]
sampleInterMax' (Int
dimInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
budget Int
minExp Int
maxExp

-- | Sample a random interaction given the problem dimension and 
-- the minimum and maximum exponents
sampleInter :: Int               -- ^ problem dimension
            -> Int               -- ^ minimum exponent            
            -> Int               -- ^ maximum exponent
            -> Rnd Interaction   -- ^ Random interaction generator
sampleInter :: Int -> Int -> Int -> Rnd Interaction
sampleInter Int
dim Int
minExp Int
maxExp = do [(Int, Int)]
es <- Int -> Int -> Int -> Rnd [(Int, Int)]
sampleInter' Int
dim Int
minExp Int
maxExp
                                   Interaction -> Rnd Interaction
forall (m :: * -> *) a. Monad m => a -> m a
return (Interaction -> Rnd Interaction) -> Interaction -> Rnd Interaction
forall a b. (a -> b) -> a -> b
$ [(Int, Int)] -> Interaction
forall a. [(Int, a)] -> IntMap a
M.fromList [(Int, Int)]
es

sampleInter' :: Int -> Int -> Int -> Rnd [(Int, Int)]
sampleInter' :: Int -> Int -> Int -> Rnd [(Int, Int)]
sampleInter' Int
0   Int
_      Int
_      = [(Int, Int)] -> Rnd [(Int, Int)]
forall (m :: * -> *) a. Monad m => a -> m a
return []
sampleInter' Int
dim Int
minExp Int
maxExp = do Int
e  <- Int -> Int -> Rnd Int
sampleRng Int
minExp Int
maxExp
                                    [(Int, Int)]
es <- Int -> Int -> Int -> Rnd [(Int, Int)]
sampleInter' (Int
dimInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
minExp Int
maxExp
                                    if Int
eInt -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
0
                                    then [(Int, Int)] -> Rnd [(Int, Int)]
forall (m :: * -> *) a. Monad m => a -> m a
return [(Int, Int)]
es
                                    else [(Int, Int)] -> Rnd [(Int, Int)]
forall (m :: * -> *) a. Monad m => a -> m a
return ((Int
dimInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1,Int
e)(Int, Int) -> [(Int, Int)] -> [(Int, Int)]
forall a. a -> [a] -> [a]
:[(Int, Int)]
es)


-- | Samples a random transformation function from a provided list of functions
sampleTrans :: [Transformation]      -- ^ choices of transformation functions
            -> Rnd Transformation  -- ^ Random generator
sampleTrans :: [Transformation] -> Rnd Transformation
sampleTrans = [Transformation] -> Rnd Transformation
forall a. [a] -> Rnd a
sampleFromList 

-- | Samples a random term using a random transformation and a random interaction generators.
sampleTerm :: Rnd Transformation  -- ^ random transformation function
           -> Rnd Interaction     -- ^ random interaction function
           -> Rnd Term            -- ^ Random generator
sampleTerm :: Rnd Transformation -> Rnd Interaction -> Rnd Term
sampleTerm Rnd Transformation
rndTrans Rnd Interaction
rndInter = do Transformation
t <- Rnd Transformation
rndTrans
                                  Transformation -> Interaction -> Term
Term Transformation
t (Interaction -> Term) -> Rnd Interaction -> Rnd Term
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Rnd Interaction
rndInter

-- | Create a random expression with exactly n terms
sampleExpr :: Rnd Term         -- ^ random term function
           -> Int                  -- ^ number of terms
           -> Rnd Expr         -- ^ Random generator
sampleExpr :: Rnd Term -> Int -> Rnd Expr
sampleExpr Rnd Term
_ Int
0            = Expr -> Rnd Expr
forall (m :: * -> *) a. Monad m => a -> m a
return []
sampleExpr Rnd Term
rndTerm Int
nTerms = do Term
t <- Rnd Term
rndTerm
                               Expr
e <- Rnd Term -> Int -> Rnd Expr
sampleExpr Rnd Term
rndTerm (Int
nTermsInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                               Expr -> Rnd Expr
forall (m :: * -> *) a. Monad m => a -> m a
return (Expr -> Rnd Expr) -> Expr -> Rnd Expr
forall a b. (a -> b) -> a -> b
$ Term
t Term -> Expr -> Expr
forall a. a -> [a] -> [a]
: Expr
e

-- | Create a random population of n expressions with varying number of terms
samplePop :: Int                   -- population size
          -> Int                   -- max number of terms
          -> (Int -> Rnd Expr) -- random expression generator
          -> Rnd [Expr]
samplePop :: Int -> Int -> (Int -> Rnd Expr) -> Rnd [Expr]
samplePop Int
nPop Int
maxNTerms Int -> Rnd Expr
rndExpr = do Int
n  <- Int -> Int -> Rnd Int
sampleRng Int
1 Int
maxNTerms
                                      Expr
e  <- Int -> Rnd Expr
rndExpr Int
n
                                      [Expr]
es <- Int -> Int -> (Int -> Rnd Expr) -> Rnd [Expr]
samplePop (Int
nPopInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1) Int
maxNTerms Int -> Rnd Expr
rndExpr
                                      [Expr] -> Rnd [Expr]
forall (m :: * -> *) a. Monad m => a -> m a
return ([Expr] -> Rnd [Expr]) -> [Expr] -> Rnd [Expr]
forall a b. (a -> b) -> a -> b
$ Expr
e Expr -> [Expr] -> [Expr]
forall a. a -> [a] -> [a]
: [Expr]
es

-- * Utility random functions

-- | sample from [0,n]
sampleTo :: Int -> Rnd Int
sampleTo :: Int -> Rnd Int
sampleTo Int
n | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0     = [Char] -> Rnd Int
forall a. HasCallStack => [Char] -> a
error [Char]
"Invalid number"
           | Bool
otherwise = Int -> Int -> Rnd Int
sampleRng Int
0 Int
n

-- | Sample from a range of integers
sampleRng :: Int -> Int -> Rnd Int
sampleRng :: Int -> Int -> Rnd Int
sampleRng Int
x Int
y = (StdGen -> (Int, StdGen)) -> Rnd Int
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((StdGen -> (Int, StdGen)) -> Rnd Int)
-> (StdGen -> (Int, StdGen)) -> Rnd Int
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> StdGen -> (Int, StdGen)
forall g a. (RandomGen g, UniformRange a) => (a, a) -> g -> (a, g)
uniformR (Int
x, Int
y)

-- | Sample from a range of integers excluding 0
sampleNZRng :: Int -> Int -> Rnd Int
sampleNZRng :: Int -> Int -> Rnd Int
sampleNZRng Int
x Int
y = do Int
z <- Int -> Int -> Rnd Int
sampleRng Int
x Int
y
                     if Int
z Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                     then Int -> Int -> Rnd Int
sampleNZRng Int
x Int
y
                     else Int -> Rnd Int
forall (m :: * -> *) a. Monad m => a -> m a
return Int
z

-- | Sample a random element from a list
sampleFromList :: [a] -> Rnd a
sampleFromList :: [a] -> Rnd a
sampleFromList [a]
xs = do let n :: Int
n = [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
xs
                       Int
i <- Int -> Rnd Int
sampleTo (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                       a -> Rnd a
forall (m :: * -> *) a. Monad m => a -> m a
return ([a]
xs [a] -> Int -> a
forall a. [a] -> Int -> a
!! Int
i)

-- | Toss a coin to your Witcher
toss :: Rnd Bool
toss :: Rnd Bool
toss = (StdGen -> (Bool, StdGen)) -> Rnd Bool
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state StdGen -> (Bool, StdGen)
forall a g. (Random a, RandomGen g) => g -> (a, g)
random