module Main (main) where

import Data.Char
import Data.List
import System.Random
import Control.Monad
import Control.Monad.State

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

-- Sum the "difference" of two strings' chars.
strDiff :: String -> Int
strDiff inp = sum $ zipWith (\a b -> if a == b then 0 else 1) targetStr inp

-- Generate a string of random chars
-- (in the range 32d - 127d)
rndString :: State StdGen String
rndString = replicateM targetLen rndChar

-- Generate a random char (in the range 32d - 127d)
rndChar :: State StdGen Char
rndChar = getRandom (randomR (' ', '~'))

-- Replace the input character with a random
-- new character with a probability of (1/length of target)
rndMutateChar :: Char -> State StdGen Char
rndMutateChar input = do
        roll  <- getRandom (randomR (0,1))
        if (roll :: Double) <= replaceProbability then rndChar else return input

-- Randomly mutate the chars of a String
rndMutateStr :: String -> State StdGen String
rndMutateStr = mapM rndMutateChar

-- Take 2 strings and return a new string that is
-- created by randomly picking each char from either
-- string       
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)

-- Generate a list (of the given length) of mutually unique random elements with
-- in the given bounds, using the given generator.
uniqueRandoms :: (Eq a , Random a) => Int -> (a, a) -> State StdGen [a]
uniqueRandoms = go []
    where go acc targetLen bounds
                    | length acc == targetLen = return acc
                    | otherwise = do
                                    randVal <- getRandom (randomR bounds)
                                    go (if randVal `notElem` acc then randVal : acc else acc) 
                                        targetLen bounds

choose a b = fst $ if snd a <= snd b then a else b
chooseR a b = if snd (snd a) > snd (snd b) then fst a else fst b

withoutCrossover :: [(String, Int)] -> Int -> State StdGen (String, Int)
withoutCrossover population count = do
        -- pick two random individuals
        [i1,i2] <- liftM (map (population !!)) (uniqueRandoms 2 (0, populationSize - 1))
        -- If we've not evolved the target string yet, carry on
        case find ((== 0) . snd) [i1,i2] of
            Just val -> return (fst val, count)
            Nothing   -> do
                newChildStr <- rndMutateStr (choose i1 i2)
                [r1,r2] <- liftM (map (\i ->(i,population !!i))) (uniqueRandoms 2 (0, populationSize - 1))
                let (start,_:end) = splitAt (chooseR r1 r2) population
                let newPopulation = start ++ ((newChildStr, strDiff newChildStr) : end)
                withoutCrossover newPopulation (count + 1)

withCrossover :: [(String, Int)] -> Int -> State StdGen (String,Int)
withCrossover population count = do
        -- pick four random individuals
        [i1,i2,i3,i4] <- liftM (map (population !!)) (uniqueRandoms 4 (0, populationSize - 1))
        -- If we've not evolved the target string yet, carry on
        case find ((== 0) . snd) [i1,i2,i3,i4] of
            Just val -> return (fst val, count)
            Nothing   -> do
                -- Pick the fitter two of two pairs of individuals to be parents
                let parentA = choose i1 i2
                let parentB = choose i3 i4
                -- create the "crossover" of the two parents
                crossoverStr <- crossoverStrings parentA parentB
                newChildStr <- rndMutateStr crossoverStr
                [r1,r2] <- liftM (map (\i-> (i,population !!i))) (uniqueRandoms 2 (0, populationSize - 1))
                let (start,_:end) = splitAt (chooseR r1 r2) population
                let newPopulation = start ++ ((newChildStr, strDiff newChildStr):end)
                withCrossover newPopulation (count + 1)

main = do
        -- Create a list of random individuals (just strings)
        gen <- newStdGen
        let individuals = replicate populationSize (fst (runState rndString gen))
        let individualFitnesses = map (\individual -> (individual, strDiff individual)) individuals
        -- Now evolve, first without and then with crossover
        newGen <- newStdGen
        print $ fst $ runState (withoutCrossover individualFitnesses populationSize) newGen
        print $ fst $ runState (withCrossover individualFitnesses populationSize) newGen