{-# LANGUAGE BangPatterns, Rank2Types #-}
{-# LANGUAGE GADTs #-}
{- |

Helper functions to run genetic algorithms and control iterations.

-}

module Moo.GeneticAlgorithm.Run (
  -- * Running algorithm
    runGA
  , runIO
  , nextGeneration
  , nextSteadyState
  , makeStoppable
  -- * Iteration control
  , loop, loopWithLog, loopIO
  , Cond(..), LogHook(..), IOHook(..)
) where

import Moo.GeneticAlgorithm.Random
import Moo.GeneticAlgorithm.Selection (bestFirst)
import Moo.GeneticAlgorithm.Types
import Moo.GeneticAlgorithm.StopCondition
import Moo.GeneticAlgorithm.Utilities (doCrossovers, doNCrossovers)

import Data.Monoid (Monoid, mempty, mappend)
import Data.Time.Clock.POSIX (getPOSIXTime)
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Control.Monad (liftM, when)

-- | Helper function to run the entire algorithm in the 'Rand' monad.
-- It takes care of generating a new random number generator.
runGA :: Rand [Genome a]             -- ^ function to create initial population
      -> ([Genome a] -> Rand b)       -- ^ genetic algorithm, see also 'loop' and 'loopWithLog'
      -> IO b                        -- ^ final population
runGA initialize ga = do
  rng <- newPureMT
  let (genomes0, rng') = runRand initialize rng
  return $ evalRand (ga genomes0) rng'

-- | Helper function to run the entire algorithm in the 'IO' monad.
runIO :: Rand [Genome a]                  -- ^ function to create initial population
      -> (IORef PureMT -> [Genome a] -> IO (Population a))
                                          -- ^ genetic algorithm, see also 'loopIO'
      -> IO (Population a)                -- ^ final population
runIO initialize gaIO = do
  rng <- newPureMT
  let (genomes0, rng') = runRand initialize rng
  rngref <- newIORef rng'
  gaIO rngref genomes0

-- | Construct a single step of the genetic algorithm.
--
-- See "Moo.GeneticAlgorithm.Binary" and "Moo.GeneticAlgorithm.Continuous"
-- for the building blocks of the algorithm.
--
nextGeneration
    :: (ObjectiveFunction objectivefn a)
    => ProblemType          -- ^ a type of the optimization @problem@
    -> objectivefn          -- ^ objective function
    -> SelectionOp a        -- ^ selection operator
    -> Int                  -- ^ @elite@, the number of genomes to keep intact
    -> CrossoverOp a        -- ^ crossover operator
    -> MutationOp a         -- ^ mutation operator
    -> StepGA Rand a
nextGeneration problem objective selectOp elite xoverOp mutationOp =
  makeStoppable objective $ \pop -> do
    genomes' <- liftM (map takeGenome) $ withElite problem elite selectOp pop
    let top = take elite genomes'
    let rest = drop elite genomes'
    genomes' <- shuffle rest         -- just in case if @selectOp@ preserves order
    genomes' <- doCrossovers genomes' xoverOp
    genomes' <- mapM mutationOp genomes'
    return $ evalObjective objective (top ++ genomes')


-- | Construct a single step of the incremental (steady-steate) genetic algorithm.
-- Exactly @n@ worst solutions are replaced with newly born children.
--
-- See "Moo.GeneticAlgorithm.Binary" and "Moo.GeneticAlgorithm.Continuous"
-- for the building blocks of the algorithm.
--
nextSteadyState
    :: (ObjectiveFunction objectivefn a)
    => Int                  -- ^ @n@, number of worst solutions to replace
    -> ProblemType          -- ^ a type of the optimization @problem@
    -> objectivefn          -- ^ objective function
    -> SelectionOp a        -- ^ selection operator
    -> CrossoverOp a        -- ^ crossover operator
    -> MutationOp a         -- ^ mutation operator
    -> StepGA Rand a
nextSteadyState n problem objective selectOp crossoverOp mutationOp =
    makeStoppable objective $ \pop -> do
      let popsize = length pop
      parents <- liftM (map takeGenome) (selectOp pop)
      children <- mapM mutationOp =<< doNCrossovers n parents crossoverOp
      let sortedPop = bestFirst problem pop
      let cpop = evalObjective objective children
      return . take popsize $ cpop ++ sortedPop


-- | Wrap a population transformation with pre- and post-conditions
-- to indicate the end of simulation.
--
-- Use this function to define custom replacement strategies
-- in addition to 'nextGeneration' and 'nextSteadyState'.
makeStoppable
    :: (ObjectiveFunction objectivefn a, Monad m)
    => objectivefn
    -> (Population a -> m (Population a))  -- single step
    -> StepGA m a
makeStoppable objective onestep stop input = do
  let pop = either (evalObjective objective) id input
  if isGenomes input && evalCond stop pop
     then return $ StopGA pop   -- stop before the first iteration
     else do
       newpop <- onestep pop
       return $ if evalCond stop newpop
                then StopGA newpop
                else ContinueGA newpop
  where
    isGenomes (Left _) = True
    isGenomes (Right _) = False


-- | Select @n@ best genomes, then select more genomes from the
-- /entire/ population (elite genomes inclusive). Elite genomes will
-- be the first in the list.
withElite :: ProblemType -> Int -> SelectionOp a -> SelectionOp a
withElite problem n select = \population -> do
  let elite = take n . eliteGenomes $ population
  selected <- select population
  return (elite ++ selected)
  where
    eliteGenomes = bestFirst problem

-- | Run strict iterations of the genetic algorithm defined by @step@.
-- Return the result of the last step.  Usually only the first two
-- arguments are given, and the result is passed to 'runGA'.
{-# INLINE loop #-}
loop :: (Monad m)
     => Cond a
     -- ^ termination condition @cond@
     -> StepGA m a
     -- ^ @step@ function to produce the next generation
     -> [Genome a]
     -- ^ initial population
     -> m (Population a)
      -- ^ final population
loop cond step genomes0 = go cond (Left genomes0)
  where
    go cond !x = do
       x' <- step cond x
       case x' of
         (StopGA pop) -> return pop
         (ContinueGA pop) -> go (updateCond pop cond) (Right pop)

-- | GA iteration interleaved with the same-monad logging hooks.
-- Usually only the first three arguments are given, and the result is
-- passed to 'runGA'.
{-# INLINE loopWithLog #-}
loopWithLog :: (Monad m, Monoid w)
     => LogHook a m w
     -- ^ periodic logging action
     -> Cond a
     -- ^ termination condition @cond@
     -> StepGA m a
     -- ^ @step@ function to produce the next generation
     -> [Genome a]
     -- ^ initial population
     -> m (Population a, w)
     -- ^ final population
loopWithLog hook cond step genomes0 = go cond 0 mempty (Left genomes0)
  where
    go cond !i !w !x = do
      x' <- step cond x
      case x' of
        (StopGA pop) -> return (pop, w)
        (ContinueGA pop) -> do
                         let w' = mappend w (runHook i pop hook)
                         let cond' = updateCond pop cond
                         go cond' (i+1) w' (Right pop)

    runHook !i !x (WriteEvery n write)
        | (rem i n) == 0 = write i x
        | otherwise      = mempty


-- | GA iteration interleaved with IO (for logging or saving the
-- intermediate results); it takes and returns the updated random
-- number generator via an IORef. Usually only the first three
-- arguments are given, and the result is passed to 'runIO'.
{-# INLINE loopIO #-}
loopIO
     :: [IOHook a]
     -- ^ input-output actions, special and time-dependent stop conditions
     -> Cond a
     -- ^ termination condition @cond@
     -> StepGA Rand a
     -- ^ @step@ function to produce the next generation
     -> IORef PureMT
     -- ^ reference to the random number generator
     -> [Genome a]
     -- ^ initial population @pop0@
     -> IO (Population a)
     -- ^ final population
loopIO hooks cond step rngref genomes0 = do
  rng <- readIORef rngref
  start <- realToFrac `liftM` getPOSIXTime
  (pop, rng') <- go start cond 0 rng (Left genomes0)
  writeIORef rngref rng'
  return pop
  where
    go start cond !i !rng !x = do
      stop <- (any id) `liftM` (mapM (runhook start i x) hooks)
      if (stop || either (const False) (evalCond cond) x)
         then return (asPopulation x, rng)
         else do
           let (x', rng') = runRand (step cond x) rng
           case x' of
             (StopGA pop) -> return (pop, rng')
             (ContinueGA pop) ->
                 do
                   let i' = i + 1
                   let cond' = updateCond pop cond
                   go start cond' i' rng' (Right pop)

    -- runhook returns True to terminate the loop
    runhook _ i x (DoEvery n io) = do
             when ((rem i n) == 0) (io i (asPopulation x))
             return False
    runhook _ _ _ (StopWhen iotest)  = iotest
    runhook start _ _ (TimeLimit limit)  = do
             now <- realToFrac `liftM` getPOSIXTime
             return (now >= start + limit)

    -- assign dummy objective value to a genome
    dummyObjective :: Genome a -> Phenotype a
    dummyObjective g = (g, 0.0)

    asPopulation = either (map dummyObjective) id

-- | Logging to run every @n@th iteration starting from 0 (the first parameter).
-- The logging function takes the current generation count and population.
data LogHook a m w where
    WriteEvery :: (Monad m, Monoid w)
               => Int
               -> (Int -> Population a -> w)
               -> LogHook a m w

-- | Input-output actions, interactive and time-dependent stop conditions.
data IOHook a
    = DoEvery { io'n :: Int, io'action :: (Int -> Population a -> IO ()) }
    -- ^ action to run every @n@th iteration, starting from 0;
    -- initially (at iteration 0) the objective value is zero.
    | StopWhen (IO Bool)
    -- ^ custom or interactive stop condition
    | TimeLimit { io't :: Double }
    -- ^ terminate iteration after @t@ seconds