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

data Args = Args { _symbols :: String, _numSymbols :: Int
                 , _states  :: String, _numStates :: Int }

runFlib :: Args -> (Char, String) -> Char -> ((Char, String), Char)
runFlib (Args smbs nsmbs sts _) (s, m) input = ((s',m), out) where
    (out:s':_) = drop (2 * (nsmbs * index s sts + index input smbs)) m
    index x    = head . elemIndices x

score :: Int -> Args -> String -> String -> Int
score run args flib input = length . filter id . zipWith (==) (tail input') .
    snd . mapAccumL (runFlib args) (head $ _states args,flib) $ init input'
    where input' = take (run + 1) $ cycle input

oneOf :: [a] -> IO a
oneOf xs = fmap (xs !!) $ randomRIO (0, length xs - 1)

replace :: Int -> a -> [a] -> [a]
replace i v xs = take i xs ++ v : drop (i + 1) xs

randomFlib :: Args -> IO String
randomFlib (Args smbs nsmbs sts nsts) = fmap concat $
    replicateM (nsmbs * nsts) (sequence [oneOf smbs, oneOf sts])

crossover :: Args -> String -> String -> IO String
crossover (Args _ nsmbs _ nsts) a b = do
    start <- randomRIO (0,         2 * nsmbs * nsts - 2)
    end   <- randomRIO (start + 1, 2 * nsmbs * nsts - 1)
    return $ take start a ++ take (end - start) (drop start b) ++ drop end a

mutate :: Args -> String -> IO String
mutate (Args smbs nsmbs sts nsts) flib = do
    i <- randomRIO (0, 2 * nsmbs * nsts - 1)
    c <- oneOf $ if mod i 2 == 0 then smbs else sts
    return $ replace i c flib

evolve :: String -> Int -> Float -> Int -> String -> IO ()
evolve states popSize breedChance run input =
    nextGen (0, "") =<< replicateM popSize (randomFlib args) where
    args = Args (map head symbols) (length symbols)
                states  (length . group $ sort states)
                where symbols = group $ sort input
    nextGen (top,_) _ | top == run = return ()
    nextGen best pop = do
        let scored = sort $ map (\flib -> (score run args flib input, flib)) pop
        let top = last scored
        breed <- fmap (< breedChance) $ randomRIO (0, 1)
        mix <- crossover args (snd $ head scored) (snd top)
        let newPop = (if breed then replace 0 mix else id) (map snd scored)
        mutIndex <- randomRIO (0, popSize - 1)
        mutant <- mutate args (newPop !! mutIndex)
        when (fst top > fst best) (print top)
        nextGen (max best top) $ replace mutIndex mutant newPop

main :: IO ()
main = evolve "ABCD" 10 0.3 100 "010011"