{-# LANGUAGE FlexibleContexts, NamedFieldPuns, RecordWildCards, TransformListComp, TupleSections, TypeApplications #-}
module Day24 (day24a, day24b) where
import Control.Arrow ((&&&))
import Data.Either (partitionEithers)
import Data.List (find, foldl', mapAccumL, sortOn)
import Data.Map.Strict (alter, assocs, fromList, lookup)
import Data.Maybe (catMaybes, fromMaybe, listToMaybe)
import Data.Ord (Down(Down))
import Data.Set (empty, insert, member, notMember)
import GHC.Exts (sortWith)
import Prelude hiding (lookup)
import Text.Megaparsec (MonadParsec, (<|>), between, notFollowedBy, optional, parseMaybe, sepBy, sepEndBy, some)
import Text.Megaparsec.Char (alphaNumChar, char, newline, printChar, string)
import Text.Megaparsec.Char.Lexer (decimal)
data Key s a = Key {team :: s, n :: a}
deriving (Eq, Ord)
data Unit a s = Unit
{ count :: a
, hp :: a
, weak :: [s]
, immune :: [s]
, power :: a
, type_ :: s
, initiative :: Int
}
parser :: (Integral a, MonadParsec e String m) => m [(Key String Int, Unit a String)]
parser = concat <$> sepEndBy teamP newline
where teamP = do
team <- some $ notFollowedBy (char ':' <|> newline) *> printChar
count <- char ':' *> newline *> sepEndBy unitP newline
return $ zipWith ((,) . Key team) [1..] count
unitP = do
count <- decimal
hp <- between (string " units each with ") (string " hit points") decimal
(weak, immune) <-
fmap (fromMaybe ([], [])) $ optional $ between (string " (") (char ')') $ do
(weaks, immunes) <- fmap partitionEithers $ flip sepBy (string "; ") $
(Left <$ string "weak to " <|> Right <$ string "immune to ") <*>
sepBy (some alphaNumChar) (string ", ")
return (concat weaks, concat immunes)
power <- string " with an attack that does " *> decimal
type_ <- char ' ' *> some alphaNumChar
initiative <- string " damage at initiative " *> decimal
return Unit {..}
step :: (Ord k, Ord i, Integral a, Ord a, Eq s) => [(Key k i, Unit a s)] -> [(Key k i, Unit a s)]
step units = assocs $ foldl' attack (fromList units) attacks
where attacks = map snd $ sortOn (Down . fst) $ catMaybes $ snd $ mapAccumL buildAttack empty
[ (key, effectivePower, type_, initiative)
| (key, Unit {count, power, type_, initiative}) <- units
, let effectivePower = count * power
, then sortWith by Down (effectivePower, initiative)
]
buildAttack used (src, effectivePower, srcType, srcInitiative) =
maybe (used, Nothing) ((`insert` used) &&& Just . (srcInitiative,) . (src,)) $
listToMaybe
[ dst
| (dst, Unit {count, power, weak, immune, initiative}) <- units
, team src /= team dst
, dst `notMember` used
, srcType `notElem` immune
, let damage = (if srcType `elem` weak then 2 else 1) * effectivePower
, then sortWith by Down (damage, count * power, initiative)
]
attack units' (src, dst)
| Just Unit {count, power, type_} <- lookup src units'
= alter (strike type_ $ count * power) dst units'
| otherwise = units'
strike type_ effectivePower (Just dst@Unit {count, hp, weak})
| count > killed = Just dst {count = count - killed}
where killed = (if type_ `elem` weak then 2 else 1) * effectivePower `div` hp
strike _ _ _ = Nothing
run :: (Ord k, Ord i, Integral a, Ord a, Eq s) => [(Key k i, Unit a s)] -> [(Key k i, Unit a s)]
run = run' empty
where run' seen units
| key `member` seen = units
| otherwise = run' (insert key seen) $ step units
where key = [(team, n, count) | (Key {team, n}, Unit {count}) <- units]
day24a :: String -> Maybe Int
day24a input = do
units <- parseMaybe @() (parser @Int) input
return $ sum $ count . snd <$> run units
day24b :: String -> Maybe Int
day24b input = do
units <- parseMaybe @() (parser @Int) input
let target = "Immune System"
units' <- find (all $ (== target) . team . fst)
[ run
[ (k, if team == target then unit {power = power + boost} else unit)
| (k@Key {team}, unit@Unit {power}) <- units
]
| boost <- [0..]
]
return $ sum $ count . snd <$> units'