{-# 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
import Data.SRTree.Recursion

-- * 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 = 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 :: 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 = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *). f (Fix f) -> Fix f
Fix (forall val. Function -> val -> SRTree val
Uni Function
f Fix SRTree
t)
replaceChild Fix SRTree
_         Fix SRTree
_ = forall a. Maybe a
Nothing 
{-# INLINE replaceChild #-}

-- Replace the children of a binary tree.
replaceChildren :: Fix SRTree -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceChildren :: Fix SRTree -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceChildren (Fix (Bin Op
f Fix SRTree
_ Fix SRTree
_)) Fix SRTree
l Fix SRTree
r = forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *). f (Fix f) -> Fix f
Fix (forall val. Op -> val -> val -> SRTree val
Bin Op
f Fix SRTree
l Fix SRTree
r)
replaceChildren Fix SRTree
_             Fix SRTree
_ Fix SRTree
_ = 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 = 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 [Int]
vars <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall p. HasVars p => p -> [Int]
_vars
               forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall val. Int -> SRTree val
Var forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall a. [a] -> StateT StdGen IO a
randomFrom [Int]
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 (Double, Double)
rng <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall p. HasVals p => p -> (Double, Double)
_range
                 forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall val. Double -> SRTree 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 (Double, Double)
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 (Int, Int)
rng <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall p. HasExps p => p -> (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 (f :: * -> *). f (Fix f) -> Fix f
Fix forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall val. Op -> val -> val -> SRTree val
Bin Op
Power Fix SRTree
0 forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall val. Double -> SRTree val
Const forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (Integral a, Num b) => a -> b
fromIntegral 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
randomFunction :: forall p. HasFuns p => RndTree p
randomFunction = do [Function]
funs <- forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks forall p. HasFuns p => p -> [Function]
_funs
                    Function
f <- forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall a. [a] -> StateT StdGen IO a
randomFrom [Function]
funs
                    forall (t :: (* -> *) -> * -> *) (m :: * -> *) a.
(MonadTrans t, Monad m) =>
m a -> t m a
lift forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *) a. Applicative f => a -> f a
pure forall a b. (a -> b) -> a -> b
$ forall (f :: * -> *). f (Fix f) -> Fix f
Fix (forall val. Function -> val -> SRTree val
Uni Function
f Fix SRTree
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
  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
8 :: Int)
  case Int
choice of
    Int
0 -> forall p. HasVars p => RndTree p
randomVar
    Int
1 -> forall p. HasVals p => RndTree p
randomConst
    Int
2 -> forall p. HasFuns p => RndTree p
randomFunction
    Int
3 -> forall p. HasExps p => RndTree p
randomPow
    Int
4 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall a b. (a -> b) -> a -> b
$ forall val. Op -> val -> val -> SRTree val
Bin Op
Add Fix SRTree
0 Fix SRTree
0
    Int
5 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall a b. (a -> b) -> a -> b
$ forall val. Op -> val -> val -> SRTree val
Bin Op
Sub Fix SRTree
0 Fix SRTree
0
    Int
6 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall a b. (a -> b) -> a -> b
$ forall val. Op -> val -> val -> SRTree val
Bin Op
Mul Fix SRTree
0 Fix SRTree
0
    Int
7 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall a b. (a -> b) -> a -> b
$ forall val. Op -> val -> val -> SRTree val
Bin Op
Div Fix SRTree
0 Fix SRTree
0
    Int
8 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall a b. (a -> b) -> a -> b
$ 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
  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
6 :: Int)
  case Int
choice of
    Int
0 -> forall p. HasFuns p => RndTree p
randomFunction
    Int
1 -> forall p. HasExps p => RndTree p
randomPow
    Int
2 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall a b. (a -> b) -> a -> b
$ forall val. Op -> val -> val -> SRTree val
Bin Op
Add Fix SRTree
0 Fix SRTree
0
    Int
3 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall a b. (a -> b) -> a -> b
$ forall val. Op -> val -> val -> SRTree val
Bin Op
Sub Fix SRTree
0 Fix SRTree
0
    Int
4 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall a b. (a -> b) -> a -> b
$ forall val. Op -> val -> val -> SRTree val
Bin Op
Mul Fix SRTree
0 Fix SRTree
0
    Int
5 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall a b. (a -> b) -> a -> b
$ forall val. Op -> val -> val -> SRTree val
Bin Op
Div Fix SRTree
0 Fix SRTree
0
    Int
6 -> forall (f :: * -> *) a. Applicative f => a -> f a
pure forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall (f :: * -> *). f (Fix f) -> Fix f
Fix forall a b. (a -> b) -> a -> b
$ 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.
randomTree :: HasEverything p => Int -> RndTree p
randomTree :: forall p. HasEverything p => Int -> RndTree p
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. HasVars p => RndTree p
randomVar
    else forall p. HasVals p => RndTree p
randomConst
randomTree Int
budget = do 
  Fix SRTree
node  <- forall p. HasEverything p => RndTree p
randomNode
  forall a. HasCallStack => Maybe a -> a
fromJust forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case Fix SRTree -> Int
arity Fix SRTree
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 Fix SRTree
node
    Int
1 -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceChild Fix SRTree
node forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall p. HasEverything p => Int -> RndTree p
randomTree (Int
budget forall a. Num a => a -> a -> a
- Int
1)
    Int
2 -> Fix SRTree -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceChildren Fix SRTree
node forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall p. HasEverything p => Int -> RndTree p
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 p. HasEverything p => Int -> RndTree p
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 :: HasEverything p => Int -> RndTree p
randomTreeBalanced :: forall p. HasEverything p => Int -> RndTree p
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. HasVars p => RndTree p
randomVar
    else forall p. HasVals p => RndTree p
randomConst
randomTreeBalanced Int
n = do 
  Fix SRTree
node  <- forall p. HasEverything p => RndTree p
randomNonTerminal
  forall a. HasCallStack => Maybe a -> a
fromJust forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> case Fix SRTree -> Int
arity Fix SRTree
node of
    Int
1 -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceChild Fix SRTree
node forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall p. HasEverything p => Int -> RndTree p
randomTreeBalanced (Int
n forall a. Num a => a -> a -> a
- Int
1)
    Int
2 -> Fix SRTree -> Fix SRTree -> Fix SRTree -> Maybe (Fix SRTree)
replaceChildren Fix SRTree
node forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall p. HasEverything p => Int -> RndTree p
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 p. HasEverything p => Int -> RndTree p
randomTreeBalanced (Int
n forall a. Integral a => a -> a -> a
`div` Int
2)