{-# LANGUAGE FlexibleInstances #-}
module BrainFuck (runBF, BF(..), strToBF) where

import qualified Data.Sequence as Sequence
import Data.Foldable (toList)
import qualified System.Random as Random
import qualified Evolve
import RandomState (splitG)
import qualified Control.Monad.State as State
import System.IO.Unsafe (unsafePerformIO)
import Control.Concurrent (yield)
import qualified EditDistance
import Data.List (intercalate)

data BF = Plus | Minus | ML | MR | Print | Loop [BF] deriving Eq

instance Show BF where
  show Plus = "+"
  show Minus = "-"
  show ML = "<"
  show MR = ">"
  show Print = "."
  show (Loop bf) = '[' : intercalate "" (map show bf) ++ "]"

mutateString :: Random.RandomGen g => String -> State.State g String
mutateString [] = do r <- fst . Random.random <$> splitG
                     if r < Evolve.mutateProbability
                       then (: []) . fst . Random.randomR ('+', ']') <$> splitG
                       else return []
mutateString (x : xs) = do i <- insert
                           d <- delete
                           r <- rest
                           return $ i ++ d ++ r
  where insert = do r <- fst . Random.random <$> splitG
                    if r < Evolve.mutateProbability
                      then (: []) . fst . Random.randomR ('+', ']') <$> splitG
                      else return []
        delete = do r <- fst . Random.random <$> splitG
                    if r < Evolve.mutateProbability
                      then return []
                      else return [x]
        rest = mutateString xs

instance Evolve.Evolvable [BF] where
  mutate bf = BrainFuck.strToBF <$> mutateString (intercalate "" (map show bf))
  score bf = 100 - fromIntegral (EditDistance.dist (BrainFuck.runBF bf) "Hello World!")

randomBF :: Random.RandomGen g => State.State g BF
randomBF = do r <- fst . Random.randomR (0, 5 :: Int) <$> splitG
              case r of
                0 -> return Plus
                1 -> return Minus
                2 -> return ML
                3 -> return MR
                4 -> return Print
                _ -> Loop <$> Evolve.mutate []

instance Random.Random BF where
  -- TODO: may as well have a correct implementation here
  randomR _ _ = undefined
  random g = State.runState randomBF g

extend :: a -> Int -> Sequence.Seq a -> Sequence.Seq a
extend def i s
  | i < length s = s
  | otherwise = s Sequence.>< Sequence.replicate (i - length s + 1) def

safeAdjust :: a -> (a -> a) -> Int -> Sequence.Seq a -> Sequence.Seq a
safeAdjust def f i s = Sequence.adjust f i (extend def i s)

iAdj :: Num a => (a -> a) -> Int -> Sequence.Seq a -> Sequence.Seq a
iAdj = safeAdjust 0

safeIndex :: a -> Sequence.Seq a -> Int -> a
safeIndex def s i = (extend def i s) `Sequence.index` i

iInd :: Num a => Sequence.Seq a -> Int -> a
iInd = safeIndex 0

safeEnumToChar :: Enum e => e -> Char
safeEnumToChar e = toEnum $ i `mod` (fromEnum '~' - spc + 1) + spc
  where i = fromEnum e
        spc = fromEnum ' '

runBF :: [BF] -> String
runBF bf = let (_, _, s) = foldl go (Sequence.singleton 0, 0, Sequence.empty) bf
           in toList s
  where go :: (Sequence.Seq Integer, Int, Sequence.Seq Char) -> BF
           -> (Sequence.Seq Integer, Int, Sequence.Seq Char)
        go (a, dp, s) Plus = (iAdj succ dp a, dp, s)
        go (a, dp, s) Minus = (iAdj pred dp a, dp, s)
        go (a, dp, s) ML = (a, max (pred dp) 0, s)
        go (a, dp, s) MR = (a, succ dp, s)
        go (a, dp, s) Print = (a, dp, s Sequence.|> safeEnumToChar (a `iInd` dp))
        go st@(a, dp, _) cmd@(Loop inner)
          | 0 <- a `iInd` dp = st
          | otherwise = unsafePerformIO yield `seq` go (foldl go st inner) cmd

strToBF' :: String -> ([BF], String)
strToBF' [] = ([], [])
strToBF' ('[':ss) =
 case strToBF' ss of
  (bf, ']':rest) -> let (bf', extra) = strToBF' rest
                    in (Loop bf : bf', extra)
  -- Failure case, but we tolerate it
  (bf, rest) -> ([Loop bf], rest)
strToBF' rest@(']':_) = ([], rest)
strToBF' (x:ss)
  | Just cmd <- conv x = let (bf, extra) = strToBF' ss
                         in (cmd : bf, extra)
  | otherwise = strToBF' ss

conv :: Char -> Maybe BF
conv '+' = Just Plus
conv '-' = Just Minus
conv '<' = Just ML
conv '>' = Just MR
conv '.' = Just Print
conv _ = Nothing

strToBF :: String -> [BF]
strToBF s = let (bf, _) = strToBF' s
            in bf