{-# language ConstraintKinds #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  Data.SRTree.Random 
-- Copyright   :  (c) Fabricio Olivetti 2021 - 2024
-- 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 Control.Monad.Reader (ReaderT, asks, runReaderT)
import Control.Monad.State.Strict ( MonadState(state), MonadTrans(lift), StateT )
import Data.Maybe (fromJust)
import Data.SRTree.Internal
import System.Random (Random (random, randomR), StdGen, mkStdGen)

-- * 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 -> [Int]
class HasVals p where
  _range :: p -> (Double, Double)
class HasExps p where
  _exponents :: p -> (Int, Int)
class HasFuns p where
  _funs :: p -> [Function]

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

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

instance HasVars FullParams where
  _vars :: FullParams -> [Int]
_vars (P [Int]
ixs (Double, Double)
_ (Int, Int)
_ [Function]
_) = [Int]
ixs
instance HasVals FullParams where
  _range :: FullParams -> (Double, Double)
_range (P [Int]
_ (Double, Double)
r (Int, Int)
_ [Function]
_) = (Double, Double)
r
instance HasExps FullParams where
  _exponents :: FullParams -> (Int, Int)
_exponents (P [Int]
_ (Double, Double)
_ (Int, Int)
e [Function]
_) = (Int, Int)
e
instance HasFuns FullParams where
  _funs :: FullParams -> [Function]
_funs (P [Int]
_ (Double, Double)
_ (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 = (StdGen -> (Bool, StdGen)) -> StateT StdGen IO Bool
forall a. (StdGen -> (a, StdGen)) -> StateT StdGen IO a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state StdGen -> (Bool, StdGen)
forall g. RandomGen g => g -> (Bool, g)
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 n <- (Int, Int) -> StateT StdGen IO Int
forall val.
(Ord val, Random val) =>
(val, val) -> StateT StdGen IO val
randomRange (Int
0, [a] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length [a]
funs Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
                     pure $ funs !! 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 = (StdGen -> (val, StdGen)) -> StateT StdGen IO val
forall a. (StdGen -> (a, StdGen)) -> StateT StdGen IO a
forall s (m :: * -> *) a. MonadState s m => (s -> (a, s)) -> m a
state ((val, val) -> StdGen -> (val, StdGen)
forall g. RandomGen g => (val, val) -> g -> (val, g)
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 :: Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceChild :: Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceChild (Fix (Uni Function
f Fix SRTree
_)) Fix SRTree
t = Fix SRTree -> Maybe (Fix SRTree)
forall a. a -> Maybe a
Just (Fix SRTree -> Maybe (Fix SRTree))
-> Fix SRTree -> Maybe (Fix SRTree)
forall a b. (a -> b) -> a -> b
$ SRTree (Fix SRTree) -> Fix SRTree
forall (f :: * -> *). f (Fix f) -> Fix f
Fix (Function -> Fix SRTree -> SRTree (Fix SRTree)
forall val. Function -> val -> SRTree val
Uni Function
f Fix SRTree
t)
replaceChild Fix SRTree
_         Fix SRTree
_ = Maybe (Fix SRTree)
forall a. Maybe a
Nothing 
{-# INLINE replaceChild #-}

-- Replace the children of a binary tree.
replaceFixChildren :: Fix SRTree -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceFixChildren :: Fix SRTree -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceFixChildren (Fix (Bin Op
f Fix SRTree
_ Fix SRTree
_)) Fix SRTree
l Fix SRTree
r = Fix SRTree -> Maybe (Fix SRTree)
forall a. a -> Maybe a
Just (Fix SRTree -> Maybe (Fix SRTree))
-> Fix SRTree -> Maybe (Fix SRTree)
forall a b. (a -> b) -> a -> b
$ SRTree (Fix SRTree) -> Fix SRTree
forall (f :: * -> *). f (Fix f) -> Fix f
Fix (Op -> Fix SRTree -> Fix SRTree -> SRTree (Fix SRTree)
forall val. Op -> val -> val -> SRTree val
Bin Op
f Fix SRTree
l Fix SRTree
r)
replaceFixChildren Fix SRTree
_             Fix SRTree
_ Fix SRTree
_ = Maybe (Fix SRTree)
forall a. Maybe a
Nothing
{-# INLINE replaceFixChildren #-}

-- | 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 = ReaderT p (StateT StdGen IO) (Fix SRTree)

-- | Returns a random variable, the parameter `p` must have the `HasVars` property
randomVar :: HasVars p => RndTree p
randomVar :: forall p. HasVars p => RndTree p
randomVar = do vars <- (p -> [Int]) -> ReaderT p (StateT StdGen IO) [Int]
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks p -> [Int]
forall p. HasVars p => p -> [Int]
_vars
               lift $ Fix . Var <$> randomFrom vars

-- | Returns a random constant, the parameter `p` must have the `HasConst` property
randomConst :: HasVals p => RndTree p
randomConst :: forall p. HasVals p => RndTree p
randomConst = do rng <- (p -> (Double, Double))
-> ReaderT p (StateT StdGen IO) (Double, Double)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks p -> (Double, Double)
forall p. HasVals p => p -> (Double, Double)
_range
                 lift $ Fix . Const <$> randomRange rng

-- | Returns a random integer power node, the parameter `p` must have the `HasExps` property
randomPow :: HasExps p => RndTree p
randomPow :: forall p. HasExps p => RndTree p
randomPow = do rng <- (p -> (Int, Int)) -> ReaderT p (StateT StdGen IO) (Int, Int)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks p -> (Int, Int)
forall p. HasExps p => p -> (Int, Int)
_exponents
               lift $ Fix . Bin Power 0 . Fix . Const . fromIntegral <$> randomRange rng

-- | Returns a random function, the parameter `p` must have the `HasFuns` property
randomFunction :: HasFuns p => RndTree p
randomFunction :: forall p. HasFuns p => RndTree p
randomFunction = do funs <- (p -> [Function]) -> ReaderT p (StateT StdGen IO) [Function]
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks p -> [Function]
forall p. HasFuns p => p -> [Function]
_funs
                    f <- lift $ randomFrom funs
                    lift $ pure $ Fix (Uni f 0)

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

-- | Returns a random non-terminal node, the parameter `p` must have every property.
randomNonTerminal :: HasEverything p => RndTree p
randomNonTerminal :: forall p. HasEverything p => RndTree p
randomNonTerminal = do
  choice <- StateT StdGen IO Int -> ReaderT p (StateT StdGen IO) Int
forall (m :: * -> *) a. Monad m => m a -> ReaderT p m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift (StateT StdGen IO Int -> ReaderT p (StateT StdGen IO) Int)
-> StateT StdGen IO Int -> ReaderT p (StateT StdGen IO) Int
forall a b. (a -> b) -> a -> b
$ (Int, Int) -> StateT StdGen IO Int
forall val.
(Ord val, Random val) =>
(val, val) -> StateT StdGen IO val
randomRange (Int
0, Int
6 :: Int)
  case choice of
    Int
0 -> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall p. HasFuns p => RndTree p
randomFunction
    Int
1 -> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall p. HasExps p => RndTree p
randomPow
    Int
2 -> Fix SRTree -> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall a. a -> ReaderT p (StateT StdGen IO) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Fix SRTree -> ReaderT p (StateT StdGen IO) (Fix SRTree))
-> (SRTree (Fix SRTree) -> Fix SRTree)
-> SRTree (Fix SRTree)
-> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SRTree (Fix SRTree) -> Fix SRTree
forall (f :: * -> *). f (Fix f) -> Fix f
Fix (SRTree (Fix SRTree) -> ReaderT p (StateT StdGen IO) (Fix SRTree))
-> SRTree (Fix SRTree) -> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall a b. (a -> b) -> a -> b
$ Op -> Fix SRTree -> Fix SRTree -> SRTree (Fix SRTree)
forall val. Op -> val -> val -> SRTree val
Bin Op
Add Fix SRTree
0 Fix SRTree
0
    Int
3 -> Fix SRTree -> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall a. a -> ReaderT p (StateT StdGen IO) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Fix SRTree -> ReaderT p (StateT StdGen IO) (Fix SRTree))
-> (SRTree (Fix SRTree) -> Fix SRTree)
-> SRTree (Fix SRTree)
-> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SRTree (Fix SRTree) -> Fix SRTree
forall (f :: * -> *). f (Fix f) -> Fix f
Fix (SRTree (Fix SRTree) -> ReaderT p (StateT StdGen IO) (Fix SRTree))
-> SRTree (Fix SRTree) -> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall a b. (a -> b) -> a -> b
$ Op -> Fix SRTree -> Fix SRTree -> SRTree (Fix SRTree)
forall val. Op -> val -> val -> SRTree val
Bin Op
Sub Fix SRTree
0 Fix SRTree
0
    Int
4 -> Fix SRTree -> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall a. a -> ReaderT p (StateT StdGen IO) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Fix SRTree -> ReaderT p (StateT StdGen IO) (Fix SRTree))
-> (SRTree (Fix SRTree) -> Fix SRTree)
-> SRTree (Fix SRTree)
-> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SRTree (Fix SRTree) -> Fix SRTree
forall (f :: * -> *). f (Fix f) -> Fix f
Fix (SRTree (Fix SRTree) -> ReaderT p (StateT StdGen IO) (Fix SRTree))
-> SRTree (Fix SRTree) -> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall a b. (a -> b) -> a -> b
$ Op -> Fix SRTree -> Fix SRTree -> SRTree (Fix SRTree)
forall val. Op -> val -> val -> SRTree val
Bin Op
Mul Fix SRTree
0 Fix SRTree
0
    Int
5 -> Fix SRTree -> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall a. a -> ReaderT p (StateT StdGen IO) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Fix SRTree -> ReaderT p (StateT StdGen IO) (Fix SRTree))
-> (SRTree (Fix SRTree) -> Fix SRTree)
-> SRTree (Fix SRTree)
-> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SRTree (Fix SRTree) -> Fix SRTree
forall (f :: * -> *). f (Fix f) -> Fix f
Fix (SRTree (Fix SRTree) -> ReaderT p (StateT StdGen IO) (Fix SRTree))
-> SRTree (Fix SRTree) -> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall a b. (a -> b) -> a -> b
$ Op -> Fix SRTree -> Fix SRTree -> SRTree (Fix SRTree)
forall val. Op -> val -> val -> SRTree val
Bin Op
Div Fix SRTree
0 Fix SRTree
0
    Int
6 -> Fix SRTree -> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall a. a -> ReaderT p (StateT StdGen IO) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Fix SRTree -> ReaderT p (StateT StdGen IO) (Fix SRTree))
-> (SRTree (Fix SRTree) -> Fix SRTree)
-> SRTree (Fix SRTree)
-> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. SRTree (Fix SRTree) -> Fix SRTree
forall (f :: * -> *). f (Fix f) -> Fix f
Fix (SRTree (Fix SRTree) -> ReaderT p (StateT StdGen IO) (Fix SRTree))
-> SRTree (Fix SRTree) -> ReaderT p (StateT StdGen IO) (Fix SRTree)
forall a b. (a -> b) -> a -> b
$ Op -> Fix SRTree -> Fix SRTree -> SRTree (Fix SRTree)
forall val. Op -> val -> val -> SRTree val
Bin Op
Power Fix SRTree
0 Fix SRTree
0
    
-- | Returns a random tree with a limited budget, the parameter `p` must have every property.
--
-- >>> let treeGen = runReaderT (randomTree 12) (P [0,1] (-10, 10) (2, 3) [Log, Exp])
-- >>> tree <- evalStateT treeGen (mkStdGen 52)
-- >>> showExpr tree
-- "(-2.7631152121655838 / Exp((x0 / ((x0 * -7.681722660704317) - Log(3.378309080134594)))))"
randomTree :: HasEverything p => Int -> RndTree p
randomTree :: forall p. HasEverything p => Int -> RndTree p
randomTree Int
0      = do
  coin <- StateT StdGen IO Bool -> ReaderT p (StateT StdGen IO) Bool
forall (m :: * -> *) a. Monad m => m a -> ReaderT p m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift StateT StdGen IO Bool
toss
  if coin
    then randomVar
    else randomConst
randomTree Int
budget = do 
  node  <- RndTree p
forall p. HasEverything p => RndTree p
randomNode
  fromJust <$> case arity node of
    Int
0 -> Maybe (Fix SRTree)
-> ReaderT p (StateT StdGen IO) (Maybe (Fix SRTree))
forall a. a -> ReaderT p (StateT StdGen IO) a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe (Fix SRTree)
 -> ReaderT p (StateT StdGen IO) (Maybe (Fix SRTree)))
-> Maybe (Fix SRTree)
-> ReaderT p (StateT StdGen IO) (Maybe (Fix SRTree))
forall a b. (a -> b) -> a -> b
$ Fix SRTree -> Maybe (Fix SRTree)
forall a. a -> Maybe a
Just Fix SRTree
node
    Int
1 -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceChild Fix SRTree
node (Fix SRTree -> Maybe (Fix SRTree))
-> RndTree p -> ReaderT p (StateT StdGen IO) (Maybe (Fix SRTree))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> RndTree p
forall p. HasEverything p => Int -> RndTree p
randomTree (Int
budget Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    Int
2 -> Fix SRTree -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceFixChildren Fix SRTree
node (Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree))
-> RndTree p
-> ReaderT p (StateT StdGen IO) (Fix SRTree -> Maybe (Fix SRTree))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> RndTree p
forall p. HasEverything p => Int -> RndTree p
randomTree (Int
budget Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) ReaderT p (StateT StdGen IO) (Fix SRTree -> Maybe (Fix SRTree))
-> RndTree p -> ReaderT p (StateT StdGen IO) (Maybe (Fix SRTree))
forall a b.
ReaderT p (StateT StdGen IO) (a -> b)
-> ReaderT p (StateT StdGen IO) a -> ReaderT p (StateT StdGen IO) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> RndTree p
forall p. HasEverything p => Int -> RndTree p
randomTree (Int
budget Int -> Int -> Int
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.
--
-- >>> let treeGen = runReaderT (randomTreeBalanced 10) (P [0,1] (-10, 10) (2, 3) [Log, Exp])
-- >>> tree <- evalStateT treeGen (mkStdGen 42)
-- >>> showExpr tree
-- "Exp(Log((((7.784360517385774 * x0) - (3.6412224491658223 ^ x1)) ^ ((x0 ^ -4.09764995657091) + Log(-7.710216839988497)))))"
randomTreeBalanced :: HasEverything p => Int -> RndTree p
randomTreeBalanced :: forall p. HasEverything p => Int -> RndTree p
randomTreeBalanced Int
n | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
1 = do
  coin <- StateT StdGen IO Bool -> ReaderT p (StateT StdGen IO) Bool
forall (m :: * -> *) a. Monad m => m a -> ReaderT p m a
forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift StateT StdGen IO Bool
toss
  if coin
    then randomVar
    else randomConst
randomTreeBalanced Int
n = do 
  node  <- RndTree p
forall p. HasEverything p => RndTree p
randomNonTerminal
  fromJust <$> case arity node of
    Int
1 -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceChild Fix SRTree
node (Fix SRTree -> Maybe (Fix SRTree))
-> RndTree p -> ReaderT p (StateT StdGen IO) (Maybe (Fix SRTree))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> RndTree p
forall p. HasEverything p => Int -> RndTree p
randomTreeBalanced (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
    Int
2 -> Fix SRTree -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceFixChildren Fix SRTree
node (Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree))
-> RndTree p
-> ReaderT p (StateT StdGen IO) (Fix SRTree -> Maybe (Fix SRTree))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> RndTree p
forall p. HasEverything p => Int -> RndTree p
randomTreeBalanced (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2) ReaderT p (StateT StdGen IO) (Fix SRTree -> Maybe (Fix SRTree))
-> RndTree p -> ReaderT p (StateT StdGen IO) (Maybe (Fix SRTree))
forall a b.
ReaderT p (StateT StdGen IO) (a -> b)
-> ReaderT p (StateT StdGen IO) a -> ReaderT p (StateT StdGen IO) b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Int -> RndTree p
forall p. HasEverything p => Int -> RndTree p
randomTreeBalanced (Int
n Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` Int
2)