import Control.Monad.Random
import Control.Applicative
import Control.Monad
import Control.Monad.Trans.Maybe
import Data.List
import Data.Maybe
import Data.Ord
import System.IO.Unsafe


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-}

type StratFunction = Int -> [Int] -> PlayerState -> IO Bool
maxi xs = maximumBy (comparing fst) (zip xs [0..])
replace pos newVal list = take pos list ++ newVal : drop (pos+1) list

playZombieRound strat playerIndex scores = playZombieRound' strat playerIndex scores defaultPlayer


playZombieRound' :: (RandomGen g) => StratFunction -> Int -> [Int] -> PlayerState -> Rand g Int
playZombieRound' strat index scores state =  do
    if continue 
       then do 
           newState <- playRound (Just state)
           if (isNothing newState) 
              then return (scores !! index)
              else playZombieRound' strat index scores (fromJust newState)
       else
           return $ (scores !! index) + (brains state)
    where continue = unsafePerformIO $ strat index scores state


playZombieDice playerStrats = evalRandIO $ playZombieDice' playerStrats 0 (-1) $ take (length playerStrats) $ repeat 0

playZombieDice' playerStrats playerIndex final scores = 
  if playerIndex == nPlayers 
    then playZombieDice' playerStrats 0 final scores
    else if final == playerIndex then return $ snd $ maxi scores
         else if (foldl max 0 scores) >= 13 && final == -1
              then playZombieDice' playerStrats playerIndex newFinal scores
              else do 
                  newscore <- playZombieRound (playerStrats !! playerIndex) playerIndex scores
                  let newscores = replace playerIndex newscore scores in
                    playZombieDice' playerStrats (playerIndex + 1) final newscores
  where nPlayers = length scores
        newFinal = if (playerIndex == 0) then (nPlayers - 1) else (playerIndex - 1)

humanPlayer index scores state = do
    putStrLn "\n"
    putStrLn $ "Your Score: " ++ show (scores !! index)
    putStrLn $ "Scores: " ++ show (scores) ++ " Index " ++ show index
    putStrLn $ "Your State:"
    print state
    putStrLn "Continue?"
    response <- getLine
    if response == "y" then return True else return False

getNPlayer :: Int -> StratFunction
getNPlayer n index scores state = do
    if (brains state) < n then return True else return False

getFourPlayer = getNPlayer 4
getFivePlayer = getNPlayer 5

fiftyFiftyPlayer :: Int -> StratFunction
fiftyFiftyPlayer n index scores state = 
    if brains state == 0 then return True else do
    outcomeStates <- replicateM n (evalRandIO $ playNRounds 1 state)
    if length (filter isJust outcomeStates) > ((3 * (quot n 4)) - (quot n 7)) then return True else return False

expectimaxPlayer :: Int -> Float -> StratFunction
expectimaxPlayer n limit index scores state = do
    outcomeStates <- replicateM n (evalRandIO $ playNRounds 1 state)
    let scoretotal = sum $ map (\x -> x - (brains state)) $ map brains $ catMaybes outcomeStates in
        let badscore = failScore * (length $ filter isNothing outcomeStates) in
            let expectedValue = (fromIntegral (scoretotal + badscore)) / (fromIntegral n) in do
              {-print state-}
              {-print expectedValue -}
              if expectedValue > limit then return True else return False
  where failScore = 0  - (brains state)



main = mainN 8 100000

playAndReport strats x = do
    out <- replicateM x (playZombieDice strats)
    print $ report out

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