{-# 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 20 relief
    print $ strategy puzzle 20 relief
    print $ strategy sample 10000 part2Relief
    print $ strategy puzzle 10000 part2Relief
    where
        relief = (`div` 3)
        part2Relief = (`mod` product [3, 11, 7, 2, 19, 5, 17, 13])
        state rounds = runRounds rounds . Vec.toList
        initialState = fmap (Seq.fromList . monkeyItems)
        strategy input rounds rel = product . take 2 . reverse . sort . foldl1 (zipWith (+)) . evalState (state rounds 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
    }
    ]