{-# LANGUAGE LambdaCase #-}
module Main where
import           Control.Monad.State
import           Data.Bifunctor      (first)
import           Data.List
import           Data.Vector         (Vector)
import qualified Data.Vector         as Vec
import qualified Data.Vector.Mutable as Vec (write)
import           Data.Sequence       (Seq)
import qualified Data.Sequence       as Seq
import           Control.Arrow       (returnA, (&&&), (>>>))
main :: IO ()
main = do
    print $ strategy sample relief
    print $ strategy puzzle relief
    where
        relief = (`div` 3)
        state = runRounds 20 . Vec.toList
        initialState = fmap (Seq.fromList . monkeyItems)
        strategy input rel = product . take 2 . reverse . sort . foldl1 (zipWith (+)) . evalState (state input rel) . initialState $ input
type WorryLevel = Int
type MonkeyIndex = Int
type MonkeyItems = Seq WorryLevel
type MonkeyMap = Vector MonkeyItems
type Activity = Int
type Inspect = Int -> Int
type WorryTest = Int -> Bool
type Relief = Int -> Int
type Throw = Bool -> Int
data Monkey = Monkey {
    monkeyItems       :: [Int]
    , monkeyNumber    :: Int
    , monkeyOperation :: Inspect
    , monkeyWorryTest :: WorryTest
    , monkeyThrow     :: Throw
    }
tick :: Inspect -> Relief -> WorryTest -> Throw -> WorryLevel -> (MonkeyIndex, WorryLevel)
tick inspect relief test throw = relief . inspect >>> (throw . test) &&& returnA
throw :: MonkeyIndex -> WorryLevel -> State MonkeyMap ()
throw to level = do
    monkeys <- get
    let oldList = monkeys Vec.! to
    let modified = Vec.modify (\v -> Vec.write v to (oldList Seq.|> level))
    put $ modified monkeys
turn :: Monkey -> Relief -> State MonkeyMap Activity
turn current relief = do
    monkeys <- get
    let index = monkeyNumber current
    let activity = length $ monkeys Vec.! index
    let tick' = tick (monkeyOperation current) relief (monkeyWorryTest current) (monkeyThrow current)
    let monkeyDecisions = fmap tick' $ monkeys Vec.! index
    forM_ monkeyDecisions $ uncurry throw
    turnDone index
    pure activity
turnDone :: MonkeyIndex -> State MonkeyMap ()
turnDone index = modify (\m -> m Vec.// [(index, Seq.empty)])
gameRound :: [Monkey] -> Relief -> State MonkeyMap [Activity]
gameRound monkeys relief = do
    forM monkeys $ \monkey -> do
        turn monkey relief
runRounds :: Int -> [Monkey] -> Relief -> State MonkeyMap [[Activity]]
runRounds rounds monkeys relief = do
    forM [1..rounds] $ \_ -> do
        gameRound monkeys relief
mkWorryFunc :: Int -> Int -> Bool
mkWorryFunc x = (==0) . flip mod x
mkThrowFunc :: a -> a -> Bool -> a
mkThrowFunc x y = \case
    True  -> x
    False -> y
sample :: Vector Monkey
sample = Vec.fromList
    [Monkey {
        monkeyItems = [79, 98],
        monkeyNumber = 0,
        monkeyOperation = (*19),
        monkeyWorryTest = mkWorryFunc 23,
        monkeyThrow = mkThrowFunc 2 3
    },
    Monkey {
        monkeyItems = [54, 65, 75, 74],
        monkeyNumber = 1,
        monkeyOperation = (+6),
        monkeyWorryTest = mkWorryFunc 19,
        monkeyThrow = mkThrowFunc 2 0
    },
    Monkey {
        monkeyItems = [79, 60, 97],
        monkeyNumber = 2,
        monkeyOperation = (^2),
        monkeyWorryTest = mkWorryFunc 13,
        monkeyThrow = mkThrowFunc 1 3
    },
    Monkey {
        monkeyItems = [74],
        monkeyNumber = 3,
        monkeyOperation = (+3),
        monkeyWorryTest = mkWorryFunc 17,
        monkeyThrow = mkThrowFunc 0 1
    }
  ]
puzzle :: Vector Monkey
puzzle = Vec.fromList
    [Monkey {
        monkeyItems = [56, 56, 92, 65, 71, 61, 79],
        monkeyNumber = 0,
        monkeyOperation = (*7),
        monkeyWorryTest = mkWorryFunc 3,
        monkeyThrow = mkThrowFunc 3 7
    },
    Monkey {
        monkeyItems = [61, 85],
        monkeyNumber = 1,
        monkeyOperation = (+5),
        monkeyWorryTest = mkWorryFunc 11,
        monkeyThrow = mkThrowFunc 6 4
    },
    Monkey {
        monkeyItems = [54, 96, 82, 78, 69],
        monkeyNumber = 2,
        monkeyOperation = (^2),
        monkeyWorryTest = mkWorryFunc 7,
        monkeyThrow = mkThrowFunc 0 7
    },
    Monkey {
        monkeyItems = [57, 59, 65, 95],
        monkeyNumber = 3,
        monkeyOperation = (+4),
        monkeyWorryTest = mkWorryFunc 2,
        monkeyThrow = mkThrowFunc 5 1
    },
    Monkey {
        monkeyItems = [62, 67, 80],
        monkeyNumber = 4,
        monkeyOperation = (*17),
        monkeyWorryTest = mkWorryFunc 19,
        monkeyThrow = mkThrowFunc 2 6
    },
    Monkey {
        monkeyItems = [91],
        monkeyNumber = 5,
        monkeyOperation = (+7),
        monkeyWorryTest = mkWorryFunc 5,
        monkeyThrow = mkThrowFunc 1 4
    },
    Monkey {
        monkeyItems = [79, 83, 64, 52, 77, 56, 63, 92],
        monkeyNumber = 6,
        monkeyOperation = (+6),
        monkeyWorryTest = mkWorryFunc 17,
        monkeyThrow = mkThrowFunc 2 0
    },
    Monkey {
        monkeyItems = [50, 97, 76, 96, 80, 56],
        monkeyNumber = 7,
        monkeyOperation = (+3),
        monkeyWorryTest = mkWorryFunc 13,
        monkeyThrow = mkThrowFunc 3 5
    }
    ]