{-# language ConstraintKinds #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.SRTree.Random 
-- Copyright   :  (c) Fabricio Olivetti 2021 - 2021
-- License     :  BSD3
-- Maintainer  :  fabricio.olivetti@gmail.com
-- Stability   :  experimental
-- Portability :  ConstraintKinds
--
-- Functions to generate random trees and nodes.
--
-----------------------------------------------------------------------------
module Data.SRTree.Random 
         ( HasVars
         , HasVals
         , HasFuns
         , HasEverything
         , FullParams(..)
         , RndTree
         , randomVar
         , randomConst
         , randomPow
         , randomFunction
         , randomNode
         , randomNonTerminal
         , randomTree
         , randomTreeBalanced
         )
         where

import System.Random 
import Control.Monad.State 
import Control.Monad.Reader 
import Data.Maybe (fromJust)

import Data.SRTree.Internal

-- * Class definition of properties that a certain parameter type has.
--
-- HasVars: does `p` provides a list of the variable indices?
-- HasVals: does `p` provides a range of values for the constants?
-- HasExps: does `p` provides a range for the integral exponentes?
-- HasFuns: does `p` provides a list of allowed functions?
class HasVars p where
  _vars :: p ix val -> [ix]
class HasVals p where
  _range :: p ix val -> (val, val)
class HasExps p where
  _exponents :: p ix val -> (Int, Int)
class HasFuns p where
  _funs :: p ix val -> [Function]

-- | Constraint synonym for all properties.
type HasEverything p = (HasVars p, HasVals p, HasExps p, HasFuns p)

-- | A structure with every property
data FullParams ix val = P [ix] (val, val) (Int, Int) [Function]

instance HasVars FullParams where
  _vars :: forall ix val. FullParams ix val -> [ix]
_vars (P [ix]
ixs (val, val)
_ (Int, Int)
_ [Function]
_) = [ix]
ixs
instance HasVals FullParams where
  _range :: forall ix val. FullParams ix val -> (val, val)
_range (P [ix]
_ (val, val)
r (Int, Int)
_ [Function]
_) = (val, val)
r
instance HasExps FullParams where
  _exponents :: forall ix val. FullParams ix val -> (Int, Int)
_exponents (P [ix]
_ (val, val)
_ (Int, Int)
e [Function]
_) = (Int, Int)
e
instance HasFuns FullParams where
  _funs :: forall ix val. FullParams ix val -> [Function]
_funs (P [ix]
_ (val, val)
_ (Int, Int)
_ [Function]
fs) = [Function]
fs

-- auxiliary function to sample between False and True
toss :: StateT StdGen IO Bool
toss :: StateT StdGen IO Bool
toss = forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state forall a g. (Random a, RandomGen g) => g -> (a, g)
random
{-# INLINE toss #-}

-- returns a random element of a list
randomFrom :: [a] -> StateT StdGen IO a
randomFrom :: forall a. [a] -> StateT StdGen IO a
randomFrom [a]
funs = do Int
n <- forall val.
(Ord val, Random val) =>
(val, val) -> StateT StdGen IO val
randomRange (Int
0, forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
funs forall a. Num a => a -> a -> a
- Int
1)
                     forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ [a]
funs forall a. [a] -> Int -> a
!! Int
n
{-# INLINE randomFrom #-}

-- returns a random element within a range
randomRange :: (Ord val, Random val) => (val, val) -> StateT StdGen IO val
randomRange :: forall val.
(Ord val, Random val) =>
(val, val) -> StateT StdGen IO val
randomRange (val, val)
rng = forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state (forall a g. (Random a, RandomGen g) => (a, a) -> g -> (a, g)
randomR (val, val)
rng)
{-# INLINE randomRange #-}

-- Replace the child of a unary tree.
replaceChild :: SRTree ix val -> SRTree ix val -> Maybe (SRTree ix val)
replaceChild :: forall ix val.
SRTree ix val -> SRTree ix val -> Maybe (SRTree ix val)
replaceChild (Fun Function
g SRTree ix val
_) SRTree ix val
t = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall ix val. Function -> SRTree ix val -> SRTree ix val
Fun Function
g SRTree ix val
t
replaceChild (Pow SRTree ix val
_ Int
k) SRTree ix val
t = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> Int -> SRTree ix val
Pow SRTree ix val
t Int
k
replaceChild SRTree ix val
_         SRTree ix val
_ = forall a. Maybe a
Nothing 
{-# INLINE replaceChild #-}

-- Replace the children of a binary tree.
replaceChildren :: SRTree ix val -> SRTree ix val -> SRTree ix val -> Maybe (SRTree ix val)
replaceChildren :: forall ix val.
SRTree ix val
-> SRTree ix val -> SRTree ix val -> Maybe (SRTree ix val)
replaceChildren (Add SRTree ix val
_ SRTree ix val
_) SRTree ix val
l SRTree ix val
r     = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Add SRTree ix val
l SRTree ix val
r
replaceChildren (Sub SRTree ix val
_ SRTree ix val
_) SRTree ix val
l SRTree ix val
r     = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Sub SRTree ix val
l SRTree ix val
r
replaceChildren (Mul SRTree ix val
_ SRTree ix val
_) SRTree ix val
l SRTree ix val
r     = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Mul SRTree ix val
l SRTree ix val
r
replaceChildren (Div SRTree ix val
_ SRTree ix val
_) SRTree ix val
l SRTree ix val
r     = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Div SRTree ix val
l SRTree ix val
r
replaceChildren (Power SRTree ix val
_ SRTree ix val
_) SRTree ix val
l SRTree ix val
r   = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Power SRTree ix val
l SRTree ix val
r
replaceChildren (LogBase SRTree ix val
_ SRTree ix val
_) SRTree ix val
l SRTree ix val
r = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
LogBase SRTree ix val
l SRTree ix val
r
replaceChildren SRTree ix val
_             SRTree ix val
_ SRTree ix val
_ = forall a. Maybe a
Nothing
{-# INLINE replaceChildren #-}

-- | RndTree is a Monad Transformer to generate random trees of type `SRTree ix val` 
-- given the parameters `p ix val` using the random number generator `StdGen`.
type RndTree p ix val = ReaderT (p ix val) (StateT StdGen IO) (SRTree ix val)

-- | Returns a random variable, the parameter `p` must have the `HasVars` property
randomVar :: HasVars p => RndTree p ix val
randomVar :: forall (p :: * -> * -> *) ix val. HasVars p => RndTree p ix val
randomVar = do [ix]
vars <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall (p :: * -> * -> *) ix val. HasVars p => p ix val -> [ix]
_vars
               forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall ix val. ix -> SRTree ix val
Var forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. [a] -> StateT StdGen IO a
randomFrom [ix]
vars

-- | Returns a random constant, the parameter `p` must have the `HasConst` property
randomConst :: (Ord val, Random val, HasVals p) => RndTree p ix val
randomConst :: forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasVals p) =>
RndTree p ix val
randomConst = do (val, val)
rng <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall (p :: * -> * -> *) ix val.
HasVals p =>
p ix val -> (val, val)
_range
                 forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall ix val. val -> SRTree ix val
Const forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall val.
(Ord val, Random val) =>
(val, val) -> StateT StdGen IO val
randomRange (val, val)
rng

-- | Returns a random integer power node, the parameter `p` must have the `HasExps` property
randomPow :: (Ord val, Random val, HasExps p) => RndTree p ix val
randomPow :: forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasExps p) =>
RndTree p ix val
randomPow = do (Int, Int)
rng <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall (p :: * -> * -> *) ix val.
HasExps p =>
p ix val -> (Int, Int)
_exponents
               forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> Int -> SRTree ix val
Pow forall ix val. SRTree ix val
Empty forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall val.
(Ord val, Random val) =>
(val, val) -> StateT StdGen IO val
randomRange (Int, Int)
rng

-- | Returns a random function, the parameter `p` must have the `HasFuns` property
randomFunction :: HasFuns p => RndTree p ix val
randomFunction :: forall (p :: * -> * -> *) ix val. HasFuns p => RndTree p ix val
randomFunction = do [Function]
funs <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall (p :: * -> * -> *) ix val.
HasFuns p =>
p ix val -> [Function]
_funs
                    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ (forall ix val. Function -> SRTree ix val -> SRTree ix val
`Fun` forall ix val. SRTree ix val
Empty) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. [a] -> StateT StdGen IO a
randomFrom [Function]
funs

-- | Returns a random node, the parameter `p` must have every property.
randomNode :: (Ord val, Random val, HasEverything p) => RndTree p ix val
randomNode :: forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasEverything p) =>
RndTree p ix val
randomNode = do
  Int
choice <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall val.
(Ord val, Random val) =>
(val, val) -> StateT StdGen IO val
randomRange (Int
0, Int
9 :: Int)
  case Int
choice of
    Int
0 -> forall (p :: * -> * -> *) ix val. HasVars p => RndTree p ix val
randomVar
    Int
1 -> forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasVals p) =>
RndTree p ix val
randomConst
    Int
2 -> forall (p :: * -> * -> *) ix val. HasFuns p => RndTree p ix val
randomFunction
    Int
3 -> forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasExps p) =>
RndTree p ix val
randomPow
    Int
4 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Add forall ix val. SRTree ix val
Empty forall ix val. SRTree ix val
Empty
    Int
5 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Sub forall ix val. SRTree ix val
Empty forall ix val. SRTree ix val
Empty
    Int
6 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Mul forall ix val. SRTree ix val
Empty forall ix val. SRTree ix val
Empty
    Int
7 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Div forall ix val. SRTree ix val
Empty forall ix val. SRTree ix val
Empty
    Int
8 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Power forall ix val. SRTree ix val
Empty forall ix val. SRTree ix val
Empty
    Int
9 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
LogBase forall ix val. SRTree ix val
Empty forall ix val. SRTree ix val
Empty

-- | Returns a random non-terminal node, the parameter `p` must have every property.
randomNonTerminal :: (Ord val, Random val, HasEverything p) => RndTree p ix val
randomNonTerminal :: forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasEverything p) =>
RndTree p ix val
randomNonTerminal = do
  Int
choice <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall val.
(Ord val, Random val) =>
(val, val) -> StateT StdGen IO val
randomRange (Int
0, Int
7 :: Int)
  case Int
choice of
    Int
0 -> forall (p :: * -> * -> *) ix val. HasFuns p => RndTree p ix val
randomFunction
    Int
1 -> forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasExps p) =>
RndTree p ix val
randomPow
    Int
2 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Add forall ix val. SRTree ix val
Empty forall ix val. SRTree ix val
Empty
    Int
3 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Sub forall ix val. SRTree ix val
Empty forall ix val. SRTree ix val
Empty
    Int
4 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Mul forall ix val. SRTree ix val
Empty forall ix val. SRTree ix val
Empty
    Int
5 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Div forall ix val. SRTree ix val
Empty forall ix val. SRTree ix val
Empty
    Int
6 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
Power forall ix val. SRTree ix val
Empty forall ix val. SRTree ix val
Empty
    Int
7 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall ix val. SRTree ix val -> SRTree ix val -> SRTree ix val
LogBase forall ix val. SRTree ix val
Empty forall ix val. SRTree ix val
Empty
    
-- | Returns a random tree with a limited budget, the parameter `p` must have every property.
randomTree :: (Ord val, Random val, HasEverything p) => Int -> RndTree p ix val
randomTree :: forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasEverything p) =>
Int -> RndTree p ix val
randomTree Int
0      = do
  Bool
coin <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift StateT StdGen IO Bool
toss
  if Bool
coin
    then forall (p :: * -> * -> *) ix val. HasVars p => RndTree p ix val
randomVar
    else forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasVals p) =>
RndTree p ix val
randomConst
randomTree Int
budget = do 
  SRTree ix val
node  <- forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasEverything p) =>
RndTree p ix val
randomNode
  forall a. HasCallStack => Maybe a -> a
fromJust forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case forall ix val. SRTree ix val -> Int
arity SRTree ix val
node of
    Int
0 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall a. a -> Maybe a
Just SRTree ix val
node
    Int
1 -> forall ix val.
SRTree ix val -> SRTree ix val -> Maybe (SRTree ix val)
replaceChild SRTree ix val
node forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasEverything p) =>
Int -> RndTree p ix val
randomTree (Int
budget forall a. Num a => a -> a -> a
- Int
1)
    Int
2 -> forall ix val.
SRTree ix val
-> SRTree ix val -> SRTree ix val -> Maybe (SRTree ix val)
replaceChildren SRTree ix val
node forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasEverything p) =>
Int -> RndTree p ix val
randomTree (Int
budget forall a. Integral a => a -> a -> a
`div` Int
2) forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasEverything p) =>
Int -> RndTree p ix val
randomTree (Int
budget forall a. Integral a => a -> a -> a
`div` Int
2)
    
-- | Returns a random tree with a approximately a number `n` of nodes, the parameter `p` must have every property.
randomTreeBalanced :: (Ord val, Random val, HasEverything p) => Int -> RndTree p ix val
randomTreeBalanced :: forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasEverything p) =>
Int -> RndTree p ix val
randomTreeBalanced Int
n | Int
n forall a. Ord a => a -> a -> Bool
<= Int
1 = do
  Bool
coin <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift StateT StdGen IO Bool
toss
  if Bool
coin
    then forall (p :: * -> * -> *) ix val. HasVars p => RndTree p ix val
randomVar
    else forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasVals p) =>
RndTree p ix val
randomConst
randomTreeBalanced Int
n = do 
  SRTree ix val
node  <- forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasEverything p) =>
RndTree p ix val
randomNonTerminal
  forall a. HasCallStack => Maybe a -> a
fromJust forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case forall ix val. SRTree ix val -> Int
arity SRTree ix val
node of
    Int
1 -> forall ix val.
SRTree ix val -> SRTree ix val -> Maybe (SRTree ix val)
replaceChild SRTree ix val
node forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasEverything p) =>
Int -> RndTree p ix val
randomTreeBalanced (Int
n forall a. Num a => a -> a -> a
- Int
1)
    Int
2 -> forall ix val.
SRTree ix val
-> SRTree ix val -> SRTree ix val -> Maybe (SRTree ix val)
replaceChildren SRTree ix val
node forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasEverything p) =>
Int -> RndTree p ix val
randomTreeBalanced (Int
n forall a. Integral a => a -> a -> a
`div` Int
2) forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> forall val (p :: * -> * -> *) ix.
(Ord val, Random val, HasEverything p) =>
Int -> RndTree p ix val
randomTreeBalanced (Int
n forall a. Integral a => a -> a -> a
`div` Int
2)