{-|
Module      : IT.Mutation
Description : Mutation operators for ITEA
Copyright   : (c) Fabricio Olivetti de Franca, 2020
License     : GPL-3
Maintainer  : fabricio.olivetti@gmail.com
Stability   : experimental
Portability : POSIX

Mutation operators.
-}
module IT.Mutation where

import Data.List (nub)
import Data.IntMap.Strict as M
import Control.Monad

import IT
import IT.Algorithms
import IT.Random

-- ---------------------------------------------------------------------------


-- * Mutation builder functions


-- | Create an Add new term mutation.
--
--      Adds a random term into the expression.
--      If this term already exists in the expression,
--      it returns the original expression without modification.
--
--      You need to provide a Fitness function and a function that
--      samples a random term.
--
addTerm :: Rnd Term -> Mutation
addTerm :: Rnd Term -> Mutation
addTerm Rnd Term
rndTerm Expr
e = do Term
t <- Rnd Term
rndTerm
                       if Term
t Term -> Expr -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` Expr
e
                       then Mutation
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
e
                       else Mutation
forall (m :: * -> *) a. Monad m => a -> m a
return (Term
t Term -> Expr -> Expr
forall a. a -> [a] -> [a]
: Expr
e)

-- | Create a Drop term mutation.
--
--      Drops a random term of the expression.
--
--      You need to provide a Fitness function.
--
dropTerm :: Mutation
dropTerm :: Mutation
dropTerm Expr
e = do let n :: Int
n = Expr -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Expr
e
                Int
i <- Int -> Rnd Int
sampleTo (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                Mutation
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Expr -> Expr
removeIthTerm Int
i Expr
e)

-- | Create a Random Replace Term mutation
--
--         Replace one random strength of
--         a random term of the expression.
--         You need to provide the minimum 
--         and maximum allowed exponent
replaceTerm :: Int -> Int -> Int -> Mutation
replaceTerm :: Int -> Int -> Int -> Mutation
replaceTerm Int
dim Int
minExp Int
maxExp Expr
e = do let n :: Int
n = Expr -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Expr
e
                                     Int
i <- Int -> Rnd Int
sampleTo (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                                     let t :: Term
t  = Maybe Term -> Term
forall p. Maybe p -> p
fromJust (Maybe Term -> Term) -> Maybe Term -> Term
forall a b. (a -> b) -> a -> b
$ Int -> Expr -> Maybe Term
getIthTerm Int
i Expr
e
                                         e' :: Expr
e' = Int -> Expr -> Expr
removeIthTerm Int
i Expr
e
                                     Term
t' <- Int -> Term -> Int -> Int -> Rnd Term
rndReplaceStrength Int
dim Term
t Int
minExp Int
maxExp
                                     Mutation
forall (m :: * -> *) a. Monad m => a -> m a
return (Term
t' Term -> Expr -> Expr
forall a. a -> [a] -> [a]
: Expr
e')
  where 
    fromJust :: Maybe p -> p
fromJust (Just p
x) = p
x
    fromJust Maybe p
Nothing  = [Char] -> p
forall a. HasCallStack => [Char] -> a
error [Char]
"Couldn't get i-th term in replaceTerm"  -- this should never happen

-- | replaces a strength at random
rndReplaceStrength :: Int -> Term -> Int -> Int -> Rnd Term
rndReplaceStrength :: Int -> Term -> Int -> Int -> Rnd Term
rndReplaceStrength Int
dim (Term Transformation
tf Interaction
ps) Int
minExp Int
maxExp =
  do Int
p <- Int -> Int -> Rnd Int
sampleRng Int
minExp Int
maxExp
     Int
i <- Int -> Rnd Int
sampleTo (Int
dimInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
     let ps' :: Interaction
ps' = (Int -> Bool) -> Interaction -> Interaction
forall a. (a -> Bool) -> IntMap a -> IntMap a
M.filter (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/=Int
0) (Interaction -> Interaction) -> Interaction -> Interaction
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Interaction -> Interaction
forall a. Int -> a -> IntMap a -> IntMap a
M.insert Int
i Int
p Interaction
ps
     if Interaction -> Int
forall a. IntMap a -> Int
M.size Interaction
ps' Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
       then Term -> Rnd Term
forall (m :: * -> *) a. Monad m => a -> m a
return (Transformation -> Interaction -> Term
Term Transformation
tf Interaction
ps)
       else Term -> Rnd Term
forall (m :: * -> *) a. Monad m => a -> m a
return (Transformation -> Interaction -> Term
Term Transformation
tf Interaction
ps')

-- | replaces a random transformation function
replaceTrans :: Rnd Transformation -> Mutation
replaceTrans :: Rnd Transformation -> Mutation
replaceTrans Rnd Transformation
rndTrans Expr
e = do let n :: Int
n = Expr -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Expr
e
                             Int
i  <- Int -> Rnd Int
sampleTo (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                             Transformation
tr <- Rnd Transformation
rndTrans
                             Mutation
forall (m :: * -> *) a. Monad m => a -> m a
return (Int -> Transformation -> Expr -> Expr
forall t. (Eq t, Num t) => t -> Transformation -> Expr -> Expr
replace Int
i Transformation
tr Expr
e)
  where
    change :: Transformation -> Term -> Term
change Transformation
tr' (Term Transformation
_ Interaction
i)       = Transformation -> Interaction -> Term
Term Transformation
tr' Interaction
i
    replace :: t -> Transformation -> Expr -> Expr
replace t
0  Transformation
_  ([])     = []
    replace t
_  Transformation
_  ([])     = [Char] -> Expr
forall a. HasCallStack => [Char] -> a
error [Char]
"Empty expression in replaceTrans"
    replace t
0  Transformation
tr ((Term
t:Expr
es)) = (Transformation -> Term -> Term
change Transformation
tr Term
t Term -> Expr -> Expr
forall a. a -> [a] -> [a]
: Expr
es)
    replace t
i  Transformation
tr ((Term
t:Expr
es)) = Term
t Term -> Expr -> Expr
forall a. a -> [a] -> [a]
: t -> Transformation -> Expr -> Expr
replace (t
it -> t -> t
forall a. Num a => a -> a -> a
-t
1) Transformation
tr (Expr
es)

-- | Combine two interactions with `op` operation (use (+) or (-)
-- for positive and negative interaction)
combineInter :: (Int -> Int -> Int) -> Int -> Int -> Mutation
combineInter :: (Int -> Int -> Int) -> Int -> Int -> Mutation
combineInter Int -> Int -> Int
op Int
minExp Int
maxExp Expr
e = do let n :: Int
n = Expr -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Expr
e
                                     Int
i <- Int -> Rnd Int
sampleTo (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                                     Int
j <- Int -> Rnd Int
sampleTo (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1)
                                     let ti :: Term
ti = Maybe Term -> Term
forall p. Maybe p -> p
fromJust (Maybe Term -> Term) -> Maybe Term -> Term
forall a b. (a -> b) -> a -> b
$ Int -> Expr -> Maybe Term
getIthTerm Int
i Expr
e
                                         tj :: Term
tj = Maybe Term -> Term
forall p. Maybe p -> p
fromJust (Maybe Term -> Term) -> Maybe Term -> Term
forall a b. (a -> b) -> a -> b
$ Int -> Expr -> Maybe Term
getIthTerm Int
j Expr
e
                                         e' :: Expr
e' = Int -> Expr -> Expr
removeIthTerm Int
i Expr
e
                                         ti' :: Term
ti'= Term -> Term -> Term
combineBoth Term
ti Term
tj
                                     if  Term -> Bool
allZeros Term
ti'
                                     then Mutation
forall (m :: * -> *) a. Monad m => a -> m a
return Expr
e'
                                     else Mutation
forall (m :: * -> *) a. Monad m => a -> m a
return Mutation -> (Expr -> Expr) -> Mutation
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Expr -> Expr
forall a. Eq a => [a] -> [a]
nub Mutation -> Mutation
forall a b. (a -> b) -> a -> b
$ Term
ti' Term -> Expr -> Expr
forall a. a -> [a] -> [a]
: Expr
e'
  where
    allZeros :: Term -> Bool
allZeros (Term Transformation
_  Interaction
is) = Interaction -> Int
forall a. IntMap a -> Int
M.size Interaction
is Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
    fromJust :: Maybe p -> p
fromJust (Just p
x)     = p
x
    fromJust Maybe p
Nothing      = [Char] -> p
forall a. HasCallStack => [Char] -> a
error [Char]
"Couldn't get a term in combineInter"

    combineBoth :: Term -> Term -> Term
combineBoth (Term Transformation
tr1 Interaction
int1) (Term Transformation
_ Interaction
int2) = Transformation -> Interaction -> Term
Term Transformation
tr1 ((Int -> Bool) -> Interaction -> Interaction
forall a. (a -> Bool) -> IntMap a -> IntMap a
M.filter (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
/=Int
0) (Interaction -> Interaction) -> Interaction -> Interaction
forall a b. (a -> b) -> a -> b
$ (Int -> Int -> Int) -> Interaction -> Interaction -> Interaction
forall a. (a -> a -> a) -> IntMap a -> IntMap a -> IntMap a
M.unionWith (\Int
i1 Int
i2 -> Int -> Int
minmax (Int
i1 Int -> Int -> Int
`op` Int
i2)) Interaction
int1 Interaction
int2)
    minmax :: Int -> Int
minmax Int
x = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
maxExp (Int -> Int) -> Int -> Int
forall a b. (a -> b) -> a -> b
$ Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
minExp Int
x

-- | Positive and Negative interaction mutations
positiveInter, negativeInter :: Int -> Int -> Mutation
positiveInter :: Int -> Int -> Mutation
positiveInter = (Int -> Int -> Int) -> Int -> Int -> Mutation
combineInter Int -> Int -> Int
forall a. Num a => a -> a -> a
(+)
negativeInter :: Int -> Int -> Mutation
negativeInter = (Int -> Int -> Int) -> Int -> Int -> Mutation
combineInter (-)

-- | Apply one of the mutation functions at random
mutFun :: Int                    -- ^ Dim
       -> (Int, Int)             -- ^ minExp, maxExp
       -> (Int, Int)             -- ^ minTerms, maxTerms
       -> Rnd Term           -- ^ random term generator
       -> Rnd Transformation -- ^ random term generator
       -> Expr                 -- ^ Expression to be mutated
       -> Rnd Expr           -- ^ Random Expression generator
mutFun :: Int
-> (Int, Int)
-> (Int, Int)
-> Rnd Term
-> Rnd Transformation
-> Mutation
mutFun Int
dim (Int
minExp, Int
maxExp) (Int
minTerms, Int
maxTerms) Rnd Term
rndTerm Rnd Transformation
rndTrans Expr
e = StateT StdGen Identity (Rnd Expr) -> Rnd Expr
forall (m :: * -> *) a. Monad m => m (m a) -> m a
join ([Rnd Expr] -> StateT StdGen Identity (Rnd Expr)
forall a. [a] -> Rnd a
sampleFromList [Rnd Expr]
muts)
  where
    muts :: [Rnd Expr]
muts   = [Int -> Int -> Int -> Mutation
replaceTerm Int
dim Int
minExp Int
maxExp Expr
e
             ,Rnd Transformation -> Mutation
replaceTrans Rnd Transformation
rndTrans         Expr
e
             ,Int -> Int -> Mutation
positiveInter Int
minExp Int
maxExp   Expr
e
             ,Int -> Int -> Mutation
negativeInter Int
minExp Int
maxExp   Expr
e] [Rnd Expr] -> [Rnd Expr] -> [Rnd Expr]
forall a. [a] -> [a] -> [a]
++ [Rnd Expr]
addMut [Rnd Expr] -> [Rnd Expr] -> [Rnd Expr]
forall a. [a] -> [a] -> [a]
++ [Rnd Expr]
dropMut

    addMut :: [Rnd Expr]
addMut  = [Rnd Term -> Mutation
addTerm Rnd Term
rndTerm Expr
e | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<  Int
maxTerms]
    dropMut :: [Rnd Expr]
dropMut = [Mutation
dropTerm Expr
e        | Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>  Int
minTerms]
    len :: Int
len     = Expr -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length Expr
e