import System.Random
import Control.Monad.State
import Data.Maybe
import Data.Tuple
import Data.Sequence as Seq
import Data.Foldable

-- | Take a random element from a list, removing it from the list.
randomChoice :: RandomGen g => State (Seq a, g) (Maybe a)
randomChoice = get >>= (\(xs, g) -> if (Seq.null xs) then (return Nothing) else (let
    (r, g2) = randomR (0, Seq.length xs - 1) g
    in (put (deleteAt r xs, g2) >> (return (Just $ index xs r))
  )))

-- | Returns a list of randomly chosen values with no repeats.
randomChoices :: RandomGen g => Int -> State (Seq a, g) (Seq a)
randomChoices n = state (\si -> 
    let (maybeVals, sf) = runState (sequence $ Prelude.replicate n randomChoice) si
    in  (fromList $ catMaybes maybeVals, sf)
  )

data Bacterium = Red | Green deriving Show
type Bacteria = Seq Bacterium

reproduce :: Bacteria -> Bacteria
reproduce b = b <> b

eat :: RandomGen g => (Bacteria, g) -> (Bacteria, g)
eat (b, g) = snd $ runState (randomChoices $ Seq.length b `quot` 2) (b, g)

hour :: RandomGen g => (Bacteria, g) -> (Bacteria, g)
hour (b, g) = eat (reproduce b, g)

initialize :: Int -> Bacteria
initialize n = (Seq.replicate (n `quot` 2) Red) <> (Seq.replicate (n `quot` 2) Green)

isRed :: Bacterium -> Bool
isRed Red   = True
isRed Green = False

nRed :: Bacteria -> Int
nRed b = Seq.length $ Seq.filter isRed b

isDone :: Bacteria -> Bool
isDone b | nRed b == 0            = True
         | nRed b == Seq.length b = True
         | otherwise              = False

run :: RandomGen g => Int -> g -> ([Bacteria], g)
run n g = (fst <$> steps, snd $ last steps)
          where steps = takeWhile (not . isDone . fst) $ iterate hour (initialize n, g)

lifetime :: [Bacteria] -> Int
lifetime b = Data.Foldable.length b