module Day21 where
import Control.Monad
import Debug.Trace (traceShow, trace)
import Data.Array

data NdState s a = NdState (s -> [(a,s)])
instance Functor (NdState s) where
    fmap = liftM

instance Applicative (NdState s) where
    pure = return
    (<*>) = ap

instance Monad (NdState s) where
    return x = NdState (\s -> [(x,s)])
    (NdState f) >>= g = NdState $ \s0 ->
        let xs = f s0 in
            concatMap (\(x,s) -> let (NdState g') = g x in g' s) xs

runState :: NdState a b -> a -> [(b,a)]
runState (NdState f) = f

put :: a -> NdState a ()
put x = NdState $ const [((), x)]

get :: NdState a a
get = NdState (\s -> [(s,s)])

putMany :: [a] -> NdState a ()
putMany xs = NdState $ const $ map (\x -> ((), x)) xs

returnMany :: [b] -> NdState a b
returnMany xs = NdState (\s -> map (\x -> (x,s)) xs)

data PlayerState = PS {score :: Integer, field :: Int} deriving Show
data GameState = P1 PlayerState PlayerState | P2 PlayerState PlayerState deriving Show
-- type Game = State (GameState, DiceState)


type Game = NdState GameState

data GameResult = Win1 | Win2 deriving Show

gsTuple :: GameState -> (Bool,Int,Int, Int,Int)
gsTuple (P1 (PS s1 f1) (PS s2 f2)) = (False, fromIntegral s1, f1, fromIntegral s2, f2)
gsTuple (P2 (PS s1 f1) (PS s2 f2)) = (True,  fromIntegral s1, f1, fromIntegral s2, f2)

tupleGs ::  (Bool,Int,Int,Int,Int) -> GameState
tupleGs (False, s1, f1, s2, f2) = P1 (PS (fromIntegral s1) f1) (PS (fromIntegral s2) f2)
tupleGs (True,  s1, f1, s2, f2) = P2 (PS (fromIntegral s1) f1) (PS (fromIntegral s2) f2)

roll :: Game Int
roll = returnMany [1,2,3]

fieldNormalize :: Int -> Int
fieldNormalize x
    | x > 10 = fieldNormalize (x-10)
    | otherwise = x

roll3 :: Game Int
roll3 = do
    x <- roll
    y <- roll
    z <- roll
    return (x+y+z)

updatePS :: Int -> PlayerState -> PlayerState
updatePS x (PS score field) = PS (score + fromIntegral nf) nf where nf = fieldNormalize $ field + x

scores :: GameState -> (Integer,Integer)
scores (P1 x y) = (score x, score y)
scores (P2 x y) = (score x, score y)

move :: Game (Maybe GameResult)
move = do
    rolls <- roll3
    gs <- get
    let newgame = (case gs of
                    (P1 p1 p2) -> P2 (updatePS rolls p1) p2
                    (P2 p1 p2) -> P1 p1 (updatePS rolls p2))
    put newgame
    return $ case scores newgame of
        (x,_) | x >= 21 -> Just Win1
        (_,y) | y >= 21 -> Just Win2
        (_,_) -> Nothing

runGame :: GameState -> (Integer, Integer)
runGame gs = rgA ! (gsTuple gs) where
    rgA = array ((False,0,1,0,1), (True,30,10,30,10))
        [let key = (p,s1,f1,s2,f2) in (key, runGame' $ tupleGs key) |
        p <- [False, True], s1 <- [0..30], s2 <- [0..30], f1 <- [1..10], f2 <- [1..10]]
    runGame' gs = foldr (sumTup . go) (0,0) (runState move gs) where
        go (Nothing, gs) = rgA ! gsTuple gs
        go (Just Win1, gs) = (1,0)
        go (Just Win2, gs) = (0,1)
        sumTup (a,b) (x,y) = (a+x, b+y)


part2 :: Int -> Int -> (Integer, Integer)
part2 p1 p2 = runGame (P1 (PS 0 p1) (PS 0 p2))

main :: IO ()
main = do
    input <- map (last . words) . lines <$> getContents
    let p1pos = read (input !! 0) :: Int
    let p2pos = read (input !! 1) :: Int

    print (p1pos, p2pos)

    print $ part2 p1pos p2pos


    print "Bye"