module Problem84 where

import Data.Array
import Data.List
import Data.Ord
import System.Random

-- community chest
cc = [2, 17, 33]
-- chance
ch = [7, 22, 36]
-- railway
r = [5, 10, 15, 25, 35]
-- utility
u = [12, 28]

main :: IO ()
main = do
    gen  <- getStdGen
    gen' <- newStdGen
    let die = zip (randomRs (1, 4) gen :: [Int]) (randomRs (1, 4) gen')
    -- go to GO(0) or JAIL(10), STAY(40)
    ccCards <- shuffle ([0, 10] ++ replicate 10 40) <$> newStdGen
    -- go to GO(0), JAIL(10), C1(11), E3(24), H2(39), R1(5), next R(40), next R(40), next U(41)
    -- , back 3(-3), STAY(42)
    chCards <-
        shuffle ([0, 10, 11, 24, 39, 5, 40, 40, 41, -3] ++ replicate 6 42) <$> newStdGen
    let pos =
            scanl (\state (d1, d2) -> nextState state d1 d2) (0, 0, ccCards, chCards) die
    -- print . take 10 $ zip die pos
    let count = array (0, 39) (zip [0 .. 39] (repeat 0))
    let rolls = 100000
    putStrLn
        $ concatMap (show . fst)
        . take 3
        . sortOn (Down . snd)
        . assocs
        . foldl (\count pos -> accum (+) count [(pos, 1)]) count
        . map fst'
        $ take rolls pos

fst' (pos, _, _, _) = pos

nextState :: (Int, Int, [Int], [Int]) -> Int -> Int -> (Int, Int, [Int], [Int])
nextState (pos, dr, ccCards, chCards) d1 d2
    |
    -- Three double rolls
      dr' == 3       = (10, 0, ccCards, chCards)
    | pos' `elem` cc = (communityChest, dr', tail ccCards ++ [headCC], chCards)
    | pos' `elem` ch = chance
    |
    -- Go to Jail
      pos' == 30     = (10, dr', ccCards, chCards)
    | otherwise      = (pos', dr', ccCards, chCards)
  where
            -- go to GO(0), JAIL(10), STAY(40)
            -- [0,10,40,40..]
    communityChest = case headCC of
        40 -> pos'
        _  -> headCC
    -- go to GO(0), JAIL(10), C1(11), E3(24), H2(39), R1(5), next R(40), next R(40), next U(41), back 3(-3), STAY(42)
    -- [0,10,11,24,39,5,40,40,41,-3,42,42,..]
    chance = case headCH of
        42 -> (pos', dr', ccCards, tail chCards ++ [headCH])
        40 -> (nextPos r pos', dr', ccCards, tail chCards ++ [headCH])
        41 -> (nextPos u pos', dr', ccCards, tail chCards ++ [headCH])
        -3 -> if pos' - 3 == 33
            then
                ( case headCC of
                    40 -> pos' - 3
                    _  -> headCC
                , dr'
                , tail ccCards ++ [headCC]
                , tail chCards ++ [headCH]
                )
            else (pos' - 3, dr', ccCards, tail chCards ++ [headCH])
        _ -> (headCH, dr', ccCards, tail chCards ++ [headCH])
    pos'   = (pos + d1 + d2) `mod` 40
    dr'    = if d1 == d2 && pos /= 10 then dr + 1 else 0
    headCC = head ccCards
    headCH = head chCards

shuffle :: RandomGen g => [Int] -> g -> [Int]
shuffle xs gen = map (xs !!) (take (length xs) . nub $ randomRs (0, length xs - 1) gen)

nextPos :: [Int] -> Int -> Int
nextPos xs p = case nexts of
    [] -> head xs
    _  -> head nexts
    where nexts = dropWhile (<= p) xs
