module Main (main) where
import Data.Char
import Data.List
import System.Random
import Control.Monad
import Control.Monad.State
import Data.Function (on)
targetStr = "Methinks it is like a weasel"
targetLen = 28
populationSize = 500
replaceProbability = 1 / fromIntegral targetLen

getRandom func
  = do (val, gen') <- liftM func get
       put gen'
       return val
 
strDiff :: String -> Int
strDiff = sum . zipWith (\ a b -> if a == b then 0 else 1) targetStr
 
rndString :: State StdGen String
rndString = replicateM targetLen rndChar
 
rndChar :: State StdGen Char
rndChar = getRandom (randomR (' ', '~'))
 
rndMutateChar :: Char -> State StdGen Char
rndMutateChar input
  = do roll <- getRandom (randomR (0, 1))
       if (roll :: Double) <= replaceProbability then rndChar else
         return input
 
rndMutateStr :: String -> State StdGen String
rndMutateStr = mapM rndMutateChar
 
crossoverStrings :: String -> String -> State StdGen String
crossoverStrings str1 str2
  = mapM
      (\ x ->
         do b <- getRandom random
            return (if b then fst x else snd x))
      (zip str1 str2)
 
uniqueRandoms :: Int -> State StdGen [Int]
uniqueRandoms = go []
  where go acc targetLen
          | length acc == targetLen = return acc
          | otherwise =
            do randVal <- getRandom (randomR (0, populationSize - 1))
               go (if randVal `notElem` acc then randVal : acc else acc) targetLen

choose a b = fst (if on (<=) snd a b then a else b)
chooseR a b = fst (if on (>) (snd . snd) a b then a else b)
 
withoutCrossover ::[(String, Int)] -> Int -> State StdGen (String, Int)
withoutCrossover population count
  = do [i1, i2] <- liftM (map (population !!)) (uniqueRandoms 2)
       childStr <- rndMutateStr (choose i1 i2)
       let childDiff = strDiff childStr
       if childDiff == 0 then return (childStr, count) else
         do [r1, r2] <- liftM (map (ap (,) (population !!))) (uniqueRandoms 2)
            let (start, _ : end) = splitAt (chooseR r1 r2) population
            let newPopulation = start ++ ((childStr, childDiff) : end)
            withoutCrossover newPopulation (count + 1)
 
withCrossover ::[(String, Int)] -> Int -> State StdGen (String, Int)
withCrossover population count
  = do [i1, i2, i3, i4] <- liftM (map (population !!)) (uniqueRandoms 4)
       let parentA = choose i1 i2
       let parentB = choose i3 i4
       crossoverStr <- crossoverStrings parentA parentB
       childStr <- rndMutateStr crossoverStr
       let childDiff = strDiff childStr
       if childDiff == 0 then return (childStr, count + 1) else
         do [r1, r2] <- liftM (map (ap (,) (population !!))) (uniqueRandoms 2)
            let (start, _ : end) = splitAt (chooseR r1 r2) population
            let newPopulation = start ++ ((childStr, childDiff) : end)
            withCrossover newPopulation (count + 1)

main
  = do gen <- newStdGen
       let individuals = replicate populationSize (fst (runState rndString gen))
       let individualFitnesses = map (ap (,) strDiff) individuals
       newStdGen >>= print . fst .
         runState (withoutCrossover individualFitnesses populationSize)
       newStdGen >>= print . fst .
         runState (withCrossover individualFitnesses populationSize)