module Main (main) where

import Data.Char
import Data.List
import System.Random

targetStr = "Methinks it is like a weasel"
targetLen = 28
populationSize = 500
replaceProbability = (1 / fromIntegral targetLen)

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

-- Generate a string of random chars
-- (in the range 32d - 127d)
rndString :: IO String
rndString = do
        gen <- newStdGen
        return $ take targetLen (randomRs (' ', '~') gen)

-- Generate a random char (in the range 32d - 127d)
rndChar :: IO Char
rndChar = do
        gen <- newStdGen
        return . fst $ randomR (' ', '~') gen

-- Replace the input character with a random
-- new character with a probability of (1/length of target)
rndMutateChar :: Char -> IO Char
rndMutateChar input = do
        gen <- newStdGen
        let probability = fst $ randomR (0,1) gen :: Double
        if probability <= replaceProbability then rndChar else return input

mutateChars :: Char -> IO String -> IO String
mutateChars char str = do
        newChar <- rndMutateChar char
        s <- str
        return $ newChar : s

-- Randomly mutate the chars of a String
rndMutateStr :: String -> IO String
rndMutateStr str = foldr mutateChars (return []) str

-- Pick either character and append it to the passed string.
crossoverChar :: (Char, Char) -> IO String -> IO String
crossoverChar chars str = do
    gen <- newStdGen
    s <- str
    let probability = fst $ (randomR (0, 1) gen) :: Double
    return $ (if probability <= 0.5 then fst else snd) chars : s
       
-- Take 2 strings and return a new string that is
-- created by randomly picking each char from either
-- string       
crossoverStrings :: String -> String -> IO String
crossoverStrings parentA parentB = foldr crossoverChar (return []) $ zip parentA parentB

-- 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) => (a, a) -> StdGen -> Int -> [a] -> IO [a]
uniqueRandoms bounds gen targetLen acc = do
        let (randVal, newGen) = randomR bounds gen
        if length acc == targetLen 
          then return acc 
          else if randVal `notElem` acc
                  then uniqueRandoms bounds newGen targetLen (randVal : acc)
                  else uniqueRandoms bounds newGen targetLen acc

-- Take two individuals, unpack and pass some property to a compare
-- function, returning the property accessed using another accessor function.
chooseIndividual :: (a -> a -> Bool) -> (b -> a) -> (b -> c) -> b -> b -> c
chooseIndividual cmp access returnAccess individualA individualB = returnAccess $ if cmp (access individualA) (access individualB) 
                                                                                    then individualA 
                                                                                    else individualB

withoutCrossover :: [(String, Int)] -> Int -> IO (String, Int)
withoutCrossover population count = do
        gen <- newStdGen
        -- pick two random individuals
        indices <- uniqueRandoms (0, populationSize - 1) gen 2 []
        let individuals = map (population !!) indices
        -- If we've not evolved the target string yet, carry on
        case find ((== 0) . snd) individuals of
            Just(val) -> return (fst val, count)
            Nothing   -> do
                replaceGen <- newStdGen
                -- pick the fitter individual to be the parent 
                let parentStr = chooseIndividual (<=) snd fst (individuals !! 0) (individuals !! 1)
                newChildStr <- rndMutateStr parentStr
                replaceIndices <- uniqueRandoms (0, populationSize - 1) replaceGen 2 []
                -- Keep hold of the indices, so we can split/merge the population
                let replaceIndividuals = map (\i -> (i, population !! i)) replaceIndices
                -- chose the less fit individual to be replaced
                let replaceeIndex = chooseIndividual (>) (snd.snd) fst (replaceIndividuals !! 0) (replaceIndividuals !! 1)
                let splits = splitAt replaceeIndex population
                let newPopulation = fst splits ++ [(newChildStr, strDiff newChildStr)] ++ (tail . snd) splits
                withoutCrossover newPopulation (count + 1)

withCrossover :: [(String, Int)] -> Int -> IO (String, Int)
withCrossover population count = do
        gen <- newStdGen
        -- pick four random individuals
        indices <- uniqueRandoms (0, populationSize - 1) gen 4 []
        let individuals = map (population !!) indices
        -- If we've not evolved the target string yet, carry on
        case find ((== 0) . snd) individuals of
            Just(val) -> return (fst val, count)
            Nothing   -> do
                replaceGen <- newStdGen
                -- Pick the fitter two of two pairs of individuals to be parents
                let parentA = chooseIndividual (<=) snd fst (individuals !! 0) (individuals !! 1)
                let parentB = chooseIndividual (<=) snd fst (individuals !! 2) (individuals !! 3)
                -- create the "crossover" of the two parents
                crossoverStr <- crossoverStrings parentA parentB
                newChildStr <- rndMutateStr crossoverStr
                replaceIndices <- uniqueRandoms (0, populationSize - 1) replaceGen 2 []
                -- Keep hold of the indices, so we can split/merge the population
                let replaceIndividuals = map (\i -> (i, population !! i)) replaceIndices
                -- chose the less fit individual to be replaced
                let replaceeIndex = chooseIndividual (>) (snd.snd) fst (replaceIndividuals !! 0) (replaceIndividuals !! 1)
                let splits = splitAt replaceeIndex population
                let newPopulation = fst splits ++ [(newChildStr, strDiff newChildStr)] ++ (tail . snd) splits
                withCrossover newPopulation (count + 1)

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