content: {- Minimize Rosenbrock function using real-valued genetic algorithm.
     Optimal value x* = (1,...,1). F(x*) = 0.
  
     It is a real-values genetic algorithm. The user may choose a
     mutation and crossover operators.  This example uses hooks to save
     evolution history.
  
     To run:
  
         ghc --make rosenbrock.hs
         ./rosenbrock gm undx > output.txt
  
     To visualize the output in gnuplot:
  
         % gnuplot
         > set logscale y ; set xlabel 'generation' ;
         > plot 'output.txt' u 1:2 w l t 'median', '' u 1:3 w l t 'best' lt 3
  
  
  -}
  
  import Moo.GeneticAlgorithm.Continuous
  
  import Control.Monad
  import Data.List
  import System.Environment (getArgs)
  import System.Exit (exitWith, ExitCode(..))
  import Text.Printf (printf)
  
  rosenbrock :: [Double] -> Double
  rosenbrock xs = sum . map f $ zip xs (drop 1 xs)
    where
     f (x1, x2) = 100.0 * (x2 - x1^(2::Int))^(2::Int) + (x1 - 1)^(2::Int)
  
  nvariables = 3
  xrange = (-30.0, 30.0)
  popsize = 100
  precision = 1e-5
  maxiters = 4000 :: Int
  elitesize = 10
  
  -- Rosenbrock function is minimized
  objective :: [Double] -> Objective
  objective xs = rosenbrock xs
  
  -- selection: tournament selection
  select = tournamentSelect Minimizing 3 (popsize-elitesize)
  
  -- Gaussian mutation, mutate fraction @genomeschanged@ of the population
  gm genomeschanged =
      let p = 1.0 - (1.0 - genomeschanged)**(1.0 / fromIntegral nvariables)
          s = 0.01*(snd xrange - fst xrange)
      in  gaussianMutate p s
  
  mutationOps = [ ("gm", gm 0.33) ]
  
  -- BLX-0.5 crossover
  blxa = blendCrossover 0.5
  -- UNDX crossover
  undx = unimodalCrossoverRP
  -- SBX crossover
  sbx = simulatedBinaryCrossover 2
  
  crossoverOps = [ ("blxa", blxa), ("undx", undx), ("sbx", sbx) ]
  
  printUsage = do
    putStrLn usage
    exitWith (ExitFailure 1)
    where
    usage = intercalate " " [ "rosenbrock", mops, xops ]
    mops = intercalate "|" (map fst mutationOps)
    xops = intercalate "|" (map fst crossoverOps)
  
  logStats = WriteEvery 10 $ \iterno pop ->
               let pop' =  bestFirst Minimizing pop
                   bestobjval = takeObjectiveValue $ head pop'
                   medianobjval = takeObjectiveValue $ pop' !! (length pop' `div` 2)
               in  [(iterno, medianobjval, bestobjval)]
  
  printStats :: [(Int, Objective, Objective)] -> IO ()
  printStats stats = do
    printf "# %-10s %15s %15s\n" "generation" "median" "best"
    flip mapM_ stats $ \(iterno, median, best) ->
        printf "%12d %15.3g %15.3g\n" iterno median best
  
  geneticAlgorithm mutate crossover = do
    -- initial population
    let initialize = replicateM popsize $ replicateM nvariables (getRandomR xrange)
    let stop = IfObjective ((<= precision) . minimum) `Or` Generations maxiters
    let step = nextGeneration Minimizing objective select elitesize crossover mutate
    --
    let ga = loopWithLog logStats stop step
    runGA initialize ga
  
  
  printBest :: Population Double -> IO ()
  printBest pop = do
    let bestGenome = takeGenome . head $ bestFirst Minimizing pop
    let vals = map (\x -> printf "%.5f" x) bestGenome
    putStrLn $ "# best solution: " ++ (intercalate ", " vals)
  
  -- usage: rosenbrock mutationOperator crossoverOperator
  main = do
    args <- getArgs
    conf <- case args of
             []       -> return (lookup "gm" mutationOps, lookup "undx" crossoverOps)
             (m:x:[]) -> return (lookup m mutationOps, lookup x crossoverOps)
             _        -> printUsage
    case conf of
      (Just mutate, Just crossover) -> do
         (pop, stats) <- geneticAlgorithm mutate crossover
         printStats stats
         printBest pop
         -- exit status depends on convergence
         let bestF = takeObjectiveValue . head $ bestFirst Minimizing pop
         if (abs bestF <= precision)
            then exitWith ExitSuccess
            else do
              printf "# failed to converge: best residual=%.5g, target=%g\n" bestF precision
              exitWith (ExitFailure 2)  -- failed to find a solution
      _ -> printUsage
  