import Control.Monad.Random
import Control.Applicative
import Control.Monad
import Control.Monad.Trans.Maybe
import Data.List
import Data.Maybe

data Cup = Cup {
              green :: Int,
              yellow :: Int,
              red :: Int
            } deriving (Show)

data Outcome = Brain | Runner | Shotgun  deriving (Show)
data Color = Red | Yellow | Green deriving (Show)

data OutDie = OutDie Outcome Color deriving (Show)

data PlayerState = PlayerState {
                    brains :: Int,
                    shotguns :: Int,
                    cup :: Cup,
                    hand :: Hand
                  } deriving (Show)

type Hand = [Color]


defaultCup = Cup 6 4 3
defaultPlayer = PlayerState 0 0 defaultCup []


takeGreen (Cup g y r) = Cup (g - 1) y r
takeYellow (Cup g y r) = Cup g (y - 1) r
takeRed (Cup g y r) = Cup g y (r - 1)

die :: (RandomGen g) => Rand g Int
die = getRandomR (1,6)

redDie 1 = OutDie Shotgun Red
redDie 2 = OutDie Shotgun Red
redDie 3 = OutDie Shotgun Red
redDie 4 = OutDie Runner  Red
redDie 5 = OutDie Runner  Red
redDie 6 = OutDie Brain   Red

yellowDie 1 = OutDie Shotgun Yellow
yellowDie 2 = OutDie Shotgun Yellow
yellowDie 3 = OutDie Runner  Yellow
yellowDie 4 = OutDie Runner  Yellow
yellowDie 5 = OutDie Brain   Yellow
yellowDie 6 = OutDie Brain   Yellow

greenDie 1 = OutDie Shotgun Green
greenDie 2 = OutDie Runner Green
greenDie 3 = OutDie Runner Green
greenDie 4 = OutDie Brain Green
greenDie 5 = OutDie Brain Green
greenDie 6 = OutDie Brain Green

diceInCup (Cup x y z) = x + y + z

pickDie cup = 
  if cupN == 0 
     then pickDie defaultCup 
     else
      do
        v <- (getRandomR (1, cupN))
        if v <= green cup then return (takeGreen cup, Green)
        else if v <= (green cup) + (yellow cup) then return (takeYellow cup, Yellow)
        else return (takeRed cup, Red)
  where cupN = diceInCup cup

takeDieFromPlayerCup :: (RandomGen g) => PlayerState -> Rand g PlayerState
takeDieFromPlayerCup player =
  do 
      (newcup, color) <- pickDie (cup player)
      return $ PlayerState (brains player) (shotguns player) newcup (color : (hand player))

fillHand :: (RandomGen g) => PlayerState -> Rand g PlayerState
fillHand player = 
  if length (hand player) == 3
     then return player
     else takeDieFromPlayerCup player >>= fillHand

zombieDie :: (RandomGen g) => Color -> Rand g OutDie
zombieDie Red = redDie <$> die
zombieDie Green = greenDie <$> die
zombieDie Yellow = yellowDie <$> die

clearHand player = PlayerState (brains player) (shotguns player) (cup player) []

rollHand :: (RandomGen g) => PlayerState -> Rand g [OutDie]
rollHand player = sequence (map zombieDie (hand player))

{-playerRoll :: (RandomGen g) => PlayerState -> Rand g PlayerState-}
{-playerRoll player = do-}
  {-player <- fillHand player-}
  {-outcomes <- rollHand player-}

updateOutcome :: PlayerState -> OutDie -> PlayerState
updateOutcome player (OutDie Shotgun _) = PlayerState (brains player) (1 + (shotguns player)) (cup player) (hand player)
updateOutcome player (OutDie Brain _ ) = PlayerState (1 + (brains player)) (shotguns player) (cup player) (hand player)
updateOutcome player (OutDie Runner color) = PlayerState (brains player) (shotguns player) (cup player) (color : (hand player))

acceptOutcome :: PlayerState -> [OutDie] -> PlayerState
acceptOutcome player outcomes = foldl updateOutcome (clearHand player) outcomes

report :: [Int] -> [(Int, Int)]
report input = zip id counts
  where
    sorted = sort input
    grouped = group sorted
    id = map (\x -> x !! 0) grouped
    counts = map length grouped

playNRounds n player = foldl' (>>=) (return $ Just player) (take n $ repeat playRound)

playRound :: (RandomGen g) => (Maybe PlayerState) -> Rand g (Maybe PlayerState)
playRound Nothing = return Nothing
playRound (Just player) = do
    playerWithHand <- fillHand player
    outcome <- rollHand playerWithHand
    let newPlayer = acceptOutcome playerWithHand outcome in
      if (shotguns newPlayer) >= 3
         then return Nothing
         else return $ Just newPlayer

{-firstRound = do-}
    {-playerWithHand <- fillHand defaultPlayer-}
    {-outcome <- rollHand playerWithHand-}
    {-return $ acceptOutcome playerWithHand outcome-}

main = mainN 8 100000

mainN n x = do 
    outcomeStates <- replicateM x (evalRandIO $ playNRounds n defaultPlayer)
    print "brains"
    print $ report $ map brains $ catMaybes outcomeStates
    print "shotguns"
    print $ report $ map shotguns $ catMaybes outcomeStates
    print "successes"
    print $ length $ filter isJust outcomeStates
    print "failures"
    print $ length $ filter isNothing outcomeStates