module Days.Day11 where

import AOCUtil
import Data.Bool
import Data.Char
import Data.List
import GHC.Utils.Misc

runA :: IO ()
runA = interactF "data/day10.txt" solve

runB :: IO ()
runB = interactF "data/day10.txt" solveB

solve :: String -> String
solve _ = "change 'handleItems' for A part " -- show .  uncurry (*) . first2 . reverse . sort . map insp . runRounds 20 $ initialMonkeys

solveB :: String -> String
solveB _ = show . uncurry (*) . first2 . reverse . sort . map insp . runRounds 10000 $ initialMonkeys

factor :: Int
factor = 2 * 7 * 3 * 11 * 17 * 5 * 13 * 19

type Item = Int

data Monkey = Monkey {items :: [Item], op :: Item -> Item, target :: Item -> Int, insp :: Int}

instance Show Monkey where
  show (Monkey it _ _ is) = show it ++ "@" ++ show is

initialMonkeys :: [Monkey]
initialMonkeys =
  [ Monkey {items = [54, 53], op = (* 3), target = bool 6 2 . (== 0) . (`mod` 2), insp = 0},
    Monkey {items = [95, 88, 75, 81, 91, 67, 65, 84], op = (* 11), target = bool 4 3 . (== 0) . (`mod` 7), insp = 0},
    Monkey {items = [76, 81, 50, 93, 96, 81, 83], op = (+ 6), target = bool 1 5 . (== 0) . (`mod` 3), insp = 0},
    Monkey {items = [83, 85, 85, 63], op = (+ 4), target = bool 4 7 . (== 0) . (`mod` 11), insp = 0},
    Monkey {items = [85, 52, 64], op = (+ 8), target = bool 7 0 . (== 0) . (`mod` 17), insp = 0},
    Monkey {items = [57], op = (+ 2), target = bool 3 1 . (== 0) . (`mod` 5), insp = 0},
    Monkey {items = [60, 95, 76, 66, 91], op = (^ 2), target = bool 5 2 . (== 0) . (`mod` 13), insp = 0},
    Monkey {items = [65, 84, 76, 72, 79, 65], op = (+ 5), target = bool 0 6 . (== 0) . (`mod` 19), insp = 0}
  ]

runRounds :: Int -> [Monkey] -> [Monkey]
runRounds n = nTimes n monkeyRound

monkeyRound :: [Monkey] -> [Monkey]
monkeyRound ms = foldl monkeyCycle ms [0 .. n - 1]
  where
    n = length ms

monkeyCycle :: [Monkey] -> Int -> [Monkey]
monkeyCycle ms i = clearMonkey i . foldl (.) id (map addItem toAdd) $ ms
  where
    activeMonkey = ms !! i
    toAdd = handleItems activeMonkey

handleItems :: Monkey -> [(Int, Item)]
handleItems (Monkey [] _ _ _) = [] -- Target Monkey, Item worry level
handleItems (Monkey (i : is) op target insp) = (target newWorry, newWorry) : handleItems (Monkey is op target insp)
  where
    -- newWorry = (`div` 3) . op $ i -- for A
    newWorry = (`mod` factor) . op $ i -- for B

addItem :: (Int, Item) -> [Monkey] -> [Monkey] -- Index Item
addItem (idx, item) = zipWith (\i m -> if idx == i then m {items = items m ++ [item]} else m) [0 ..]

clearMonkey :: Int -> [Monkey] -> [Monkey]
clearMonkey idx = zipWith (\i m -> if idx == i then m {items = [], insp = insp m + length (items m)} else m) [0 ..]
