{-# Language QuasiQuotes, ImportQualifiedPost, NumericUnderscores #-}
{-|
Module      : Main
Description : Day 11 solution
Copyright   : (c) Eric Mertens, 2022
License     : ISC
Maintainer  : emertens@gmail.com

<https://adventofcode.com/2022/day/11>

This solution relies on the fact that there are no interactions between
items, so each item can be simulated separately. When a monkey throws an
item to a monkey with a larger ID that monkey will get to throw the item
in the same round, but when it throws it to an earlier monkey, that item
will not be thrown until the next round. Using this we can keep track of
when an item would stop moving.

>>> :{
:main +
  "Monkey 0:\n\
  \  Starting items: 79, 98\n\
  \  Operation: new = old * 19\n\
  \  Test: divisible by 23\n\
  \    If true: throw to monkey 2\n\
  \    If false: throw to monkey 3\n\
  \\n\
  \Monkey 1:\n\
  \  Starting items: 54, 65, 75, 74\n\
  \  Operation: new = old + 6\n\
  \  Test: divisible by 19\n\
  \    If true: throw to monkey 2\n\
  \    If false: throw to monkey 0\n\
  \\n\
  \Monkey 2:\n\
  \  Starting items: 79, 60, 97\n\
  \  Operation: new = old * old\n\
  \  Test: divisible by 13\n\
  \    If true: throw to monkey 1\n\
  \    If false: throw to monkey 3\n\
  \\n\
  \Monkey 3:\n\
  \  Starting items: 74\n\
  \  Operation: new = old + 3\n\
  \  Test: divisible by 17\n\
  \    If true: throw to monkey 0\n\
  \    If false: throw to monkey 1\n"
:}
10605
2713310158

-}
module Main where

import Data.Array ( Array, (!), array )
import Data.Foldable (toList)
import Data.List (sortBy)

import Advent (format, counts)

-- | Input file contains a list of:
--
-- * Monkey ID
-- * Starting items
-- * @+@ or @*@
-- * A literal or the @old@ variable
-- * A divisor
-- * The monkey ID when the divisor divides the value
-- * The monkey ID when the divisor does not divide the value
type Input = [(Int, [Int], Char, Maybe Int, Int, Int, Int)]

-- |
-- >>> :main
-- 151312
-- 51382025916
main :: IO ()
IO ()
main =
 do [(Int, [Int], Char, Maybe Int, Int, Int, Int)]
input <- [format|2022 11
      (Monkey %u:%n
        Starting items: %u&(, )%n
        Operation: new = old %c (old|%u)%n
        Test: divisible by %u%n
          If true: throw to monkey %u%n
          If false: throw to monkey %u%n)&%n|]

    -- It's safe to work with integers mod the lcm of all the divisor tests.
    -- This will keep the worry numbers small while preserving all the
    -- divisibility tests.
    let modulus :: Int
modulus = (Int -> Int -> Int) -> [Int] -> Int
forall a. (a -> a -> a) -> [a] -> a
forall (t :: * -> *) a. Foldable t => (a -> a -> a) -> t a -> a
foldl1 Int -> Int -> Int
forall a. Integral a => a -> a -> a
lcm [Int
d | (Int
_, [Int]
_, Char
_, Maybe Int
_, Int
d, Int
_, Int
_) <- [(Int, [Int], Char, Maybe Int, Int, Int, Int)]
input]

    Int -> IO ()
forall a. Show a => a -> IO ()
print ((Int -> Int)
-> Int -> [(Int, [Int], Char, Maybe Int, Int, Int, Int)] -> Int
solve (Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div`      Int
3)     Int
20 [(Int, [Int], Char, Maybe Int, Int, Int, Int)]
input)
    Int -> IO ()
forall a. Show a => a -> IO ()
print ((Int -> Int)
-> Int -> [(Int, [Int], Char, Maybe Int, Int, Int, Int)] -> Int
solve (Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod`Int
modulus) Int
10_000 [(Int, [Int], Char, Maybe Int, Int, Int, Int)]
input)

-- | Run the given number of monkey throwing rounds and compute monkey business.
solve ::
  (Int -> Int) {- ^ extra operation to run after each monkey updates worry -} ->
  Int          {- ^ number of rounds to run -} ->
  Input        {- ^ input file contents -} ->
  Int          {- ^ product of top 2 counts of times each monkey threw something -}
solve :: (Int -> Int)
-> Int -> [(Int, [Int], Char, Maybe Int, Int, Int, Int)] -> Int
solve Int -> Int
post Int
rounds [(Int, [Int], Char, Maybe Int, Int, Int, Int)]
input = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
product ([Int] -> Int) -> ([Int] -> [Int]) -> [Int] -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> [Int] -> [Int]
forall a. Ord a => Int -> [a] -> [a]
top Int
2 ([Int] -> [Int]) -> ([Int] -> [Int]) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Map Int Int -> [Int]
forall a. Map Int a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList (Map Int Int -> [Int]) -> ([Int] -> Map Int Int) -> [Int] -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Int] -> Map Int Int
forall (f :: * -> *) a. (Foldable f, Ord a) => f a -> Map a Int
counts ([Int] -> Int) -> [Int] -> Int
forall a b. (a -> b) -> a -> b
$
  [ Int
thrower
    | (Int
i, [Int]
startingItems, Char
_, Maybe Int
_, Int
_, Int
_, Int
_) <- [(Int, [Int], Char, Maybe Int, Int, Int, Int)]
input
    , Int
item <- [Int]
startingItems
    , Int
thrower <- Int -> Int -> Int -> [Int]
go Int
rounds Int
i Int
item
  ]
  where
    -- list of monkey ids used to find bounds for monkey array
    ids :: [Int]
ids = [Int
i | (Int
i, [Int]
_, Char
_, Maybe Int
_, Int
_, Int
_, Int
_) <- [(Int, [Int], Char, Maybe Int, Int, Int, Int)]
input]

    -- store monkey information in array for faster random access
    monkeys :: Array Int (Int -> Int, Int, Int, Int)
    monkeys :: Array Int (Int -> Int, Int, Int, Int)
monkeys = (Int, Int)
-> [(Int, (Int -> Int, Int, Int, Int))]
-> Array Int (Int -> Int, Int, Int, Int)
forall i e. Ix i => (i, i) -> [(i, e)] -> Array i e
array ([Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [Int]
ids, [Int] -> Int
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
maximum [Int]
ids) [(Int
i, (Char -> Maybe Int -> Int -> Int
eval Char
o Maybe Int
n,Int
d,Int
t,Int
f)) | (Int
i, [Int]
_, Char
o, Maybe Int
n, Int
d, Int
t, Int
f) <- [(Int, [Int], Char, Maybe Int, Int, Int, Int)]
input]

    go :: Int -> Int -> Int -> [Int]
    go :: Int -> Int -> Int -> [Int]
go Int
0 Int
_ Int
_ = []
    go Int
r Int
i Int
x =
      Int
i Int -> [Int] -> [Int]
forall a. a -> [a] -> [a]
:
      case Array Int (Int -> Int, Int, Int, Int)
monkeys Array Int (Int -> Int, Int, Int, Int)
-> Int -> (Int -> Int, Int, Int, Int)
forall i e. Ix i => Array i e -> i -> e
! Int
i of
        (Int -> Int
o,Int
d,Int
t,Int
f) -> Int -> Int -> Int -> [Int]
go Int
r' Int
i' Int
x'
          where
            x' :: Int
x' = Int -> Int
post (Int -> Int
o Int
x)
            i' :: Int
i' = if Int
x' Int -> Int -> Int
forall a. Integral a => a -> a -> a
`mod` Int
d Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Int
t else Int
f
            r' :: Int
r' = if Int
i' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
i then Int
rInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1 else Int
r

-- | Evaluate a monkey's worry update function.
eval :: Char -> Maybe Int -> Int -> Int
eval :: Char -> Maybe Int -> Int -> Int
eval Char
'*' (Just Int
n) = (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
*)
eval Char
'+' (Just Int
n) = (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+)
eval Char
'*' Maybe Int
Nothing  = (Int -> Int -> Int
forall a b. (Num a, Integral b) => a -> b -> a
^(Int
2::Int))
eval Char
'+' Maybe Int
Nothing  = (Int
2Int -> Int -> Int
forall a. Num a => a -> a -> a
*)
eval Char
op Maybe Int
_ = String -> Int -> Int
forall a. HasCallStack => String -> a
error (String
"Unexpected operation: " String -> String -> String
forall a. [a] -> [a] -> [a]
++ [Char
op])

-- | Returns the largest @n@ numbers in a list
top :: Ord a => Int -> [a] -> [a]
top :: forall a. Ord a => Int -> [a] -> [a]
top Int
n = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
take Int
n ([a] -> [a]) -> ([a] -> [a]) -> [a] -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> a -> Ordering) -> [a] -> [a]
forall a. (a -> a -> Ordering) -> [a] -> [a]
sortBy ((a -> a -> Ordering) -> a -> a -> Ordering
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> a -> Ordering
forall a. Ord a => a -> a -> Ordering
compare)