src/advent24/advent24.hs

   1 {-# LANGUAGE NegativeLiterals #-}
   2 {-# LANGUAGE FlexibleContexts #-}
   3 {-# LANGUAGE OverloadedStrings #-}
   4 {-# LANGUAGE TypeFamilies #-}
   5
   6 import Debug.Trace
   7
   8 -- import Prelude hiding ((++))
   9 import Data.Text (Text)
  10 import qualified Data.Text as T
  11 import qualified Data.Text.IO as TIO
  12
  13 import Data.Void (Void)
  14 import Text.Megaparsec hiding (State)
  15 import Text.Megaparsec.Char
  16 import qualified Text.Megaparsec.Char.Lexer as L
  17 import qualified Control.Applicative as CA
  18
  19 import Data.List hiding (group)
  20 -- import Data.Function (on)
  21 import qualified Data.Map.Strict as M
  22 import Data.Map.Strict ((!))
  23
  24
  25 data Group = Group { _units :: Int
  26                    , _hps :: Int
  27                    , _modifiers :: [Modifier]
  28                    , _damage :: Int
  29                    , _damageType :: String
  30                    , _initiative :: Int
  31                    } deriving (Eq, Show)
  32 instance Ord Group where
  33     g1 `compare` g2 = if (effectivePower g1) == (effectivePower g2)
  34                       then (_initiative g1) `compare` (_initiative g2)
  35                       else (effectivePower g1) `compare` (effectivePower g2)
  36
  37 data Modifier = Weakness [String] | Immunity [String]
  38         deriving (Eq, Show)
  39
  40 type Army = M.Map Int Group
  41 type Allocation = (Int, Int)
  42 data BattleOrder = Immune Int | Infection Int deriving (Eq, Show)
  43
  44
  45 main :: IO ()
  46 main = do
  47         text <- TIO.readFile "data/advent24.txt"
  48         let armies = successfulParse text
  49         print armies
  50         print $ part1 (fst armies) (snd armies)
  51         print $ part2 (fst armies) (snd armies)
  52
  53
  54 part1 :: Army -> Army -> Int
  55 part1 immuneArmy infectionArmy = uncurry remainingUnitCount endState
  56     where endState = battle immuneArmy infectionArmy
  57
  58
  59 part2 immuneArmy infectionArmy = (minimalBoost, part1 immuneArmy' infectionArmy)
  60     where boostUpper = findSuccessfulBoost 1 immuneArmy infectionArmy
  61           minimalBoost = boostSearch 0 boostUpper immuneArmy infectionArmy
  62           immuneArmy' = applyBoost minimalBoost immuneArmy
  63
  64
  65 findSuccessfulBoost boost immuneArmy infectionArmy =
  66     if success
  67     then boost
  68     else findSuccessfulBoost (2 * boost) immuneArmy infectionArmy
  69     where success = immuneSuccessWithBoost boost immuneArmy infectionArmy
  70
  71 boostSearch lower upper _ _ | trace ("Searching in " ++ show (lower, upper)) False = undefined
  72 boostSearch lower upper immuneArmy infectionArmy =
  73     if lower == upper
  74     then lower
  75     else boostSearch lower' upper' immuneArmy infectionArmy
  76     where boost = lower + (upper - lower) `div` 2
  77           success = immuneSuccessWithBoost boost immuneArmy infectionArmy
  78           lower' = if success then lower else boost + 1
  79           upper' = if success then boost else upper
  80
  81 immuneSuccessWithBoost :: Int -> Army -> Army -> Bool
  82 immuneSuccessWithBoost boost _ _ | trace ("Trial with boost " ++ show boost) False = undefined
  83 immuneSuccessWithBoost boost immuneArmy infectionArmy = immuneWins $ battle immuneArmy' infectionArmy
  84     where immuneArmy' = applyBoost boost immuneArmy
  85
  86
  87
  88 effectivePower :: Group -> Int
  89 effectivePower group = (_units group) * (_damage group)
  90
  91 damageCaused :: Group -> Group -> Int
  92 damageCaused attacker defender =
  93     if attackType `elem` immunities then 0
  94     else if attackType `elem` weaknesses then (2 * effectivePower attacker)
  95          else effectivePower attacker
  96     where attackType = _damageType attacker
  97           weaknesses = foldl' extractWeakness [] (_modifiers defender)
  98           immunities = foldl' extractImmunity [] (_modifiers defender)
  99
 100 extractWeakness :: [String] -> Modifier -> [String]
 101 extractWeakness currentWeaknesses (Weakness ws) = currentWeaknesses ++ ws
 102 extractWeakness currentWeaknesses (Immunity _ws) = currentWeaknesses
 103
 104 extractImmunity :: [String] -> Modifier -> [String]
 105 extractImmunity currentImmunity (Weakness _ms) = currentImmunity
 106 extractImmunity currentImmunity (Immunity ms) = currentImmunity ++ ms
 107
 108 applyDamage :: Group -> Int -> Group
 109 applyDamage group damage = group { _units = unitsRemaining }
 110     where unitsKilled = damage `div` (_hps group)
 111           unitsRemaining = maximum [0, (_units group) - unitsKilled]
 112
 113
 114 keysByEffectivePower :: Army -> [Int]
 115 keysByEffectivePower army = reverse $ sortOn (\k -> effectivePower (army!k)) (M.keys army)
 116
 117 allocateAttackers :: Army -> Army -> [Allocation]
 118 allocateAttackers attackers defenders = fst $ foldl' (allocateAttacker attackers defenders) ([], M.keys defenders) $ keysByEffectivePower attackers
 119
 120 allocateAttacker :: Army -> Army -> ([Allocation], [Int]) -> Int -> ([Allocation], [Int])
 121 -- allocateAttacker attackers defenders allocated@(assignedTargets, availableTargets) attackerKey | trace ("Allocate " ++ show attackerKey ++ "\n" ++ show allocated ++ "\n" ++ show [(t, sortTarget t) | t <- targets]) False = undefined
 122 --     where targets = reverse $ sortOn sortTarget availableTargets
 123 --           sortTarget t = ( damageCaused (attackers!attackerKey) (defenders!t)
 124 --                          , effectivePower (defenders!t)
 125 --                          , _initiative (defenders!t)
 126 --                          )
 127 allocateAttacker attackers defenders (assignedTargets, availableTargets) attackerKey =
 128     if null viableTargets
 129     then (assignedTargets, availableTargets)
 130     else (((attackerKey, target):assignedTargets), delete target availableTargets)
 131     where attacker = attackers!attackerKey
 132           viableTargets = filter (\t -> damageCaused attacker (defenders!t) > 0 ) availableTargets
 133           target = head $ reverse $ sortOn sortTarget viableTargets
 134           sortTarget t = ( damageCaused attacker (defenders!t)
 135                          , effectivePower (defenders!t)
 136                          , _initiative (defenders!t)
 137                          )
 138
 139
 140 battleOrder :: Army -> Army -> [BattleOrder]
 141 battleOrder immuneArmy infectionArmy = mergeOrders immuneIds infectIds
 142     where armyIds army = reverse $ sort [(_initiative (army!k), k) | k <- M.keys army]
 143           immuneIds = armyIds immuneArmy
 144           infectIds = armyIds infectionArmy
 145
 146 mergeOrders :: [(Int, Int)] -> [(Int, Int)] -> [BattleOrder]
 147 mergeOrders [] ids = [ Infection k | (_, k) <- ids ]
 148 mergeOrders ids [] = [ Immune k | (_, k) <- ids ]
 149 mergeOrders ((i1, k1):id1s) ((i2, k2):id2s)
 150     | i1 >= i2  = (Immune k1):(mergeOrders id1s ((i2, k2):id2s))
 151     | otherwise = (Infection k2):(mergeOrders ((i1, k1):id1s) id2s)
 152
 153 battleRound :: (Army, Army) -> (Army, Army)
 154 -- battleRound (immuneArmy, infectionArmy) | trace ("Round\n" ++ show immuneArmy ++ " " ++ show infectionArmy) False = undefined
 155 -- battleRound immuneArmy infectionArmy | trace (show (armyCount immuneArmy) ++ " " ++ show (armyCount infectionArmy)) False = undefined
 156 battleRound (immuneArmy, infectionArmy) = (pruneArmy immuneArmy'', pruneArmy infectionArmy')
 157     where immuneAllocations = allocateAttackers immuneArmy infectionArmy
 158           infectionAllocations = allocateAttackers infectionArmy immuneArmy
 159           actionOrder = battleOrder immuneArmy infectionArmy
 160           (immuneArmy', infectionArmy') =     foldl' (\ (a1, a2) order -> handleOrder order immuneAllocations infectionAllocations a1 a2)
 161                                                      (immuneArmy, infectionArmy)  actionOrder
 162           -- test for stalemate
 163           immuneArmy'' = if (immuneArmy' == immuneArmy) && (infectionArmy' == infectionArmy)
 164                          then M.empty
 165                          else immuneArmy'
 166
 167
 168
 169 handleOrder :: BattleOrder -> [Allocation] -> [Allocation] -> Army -> Army -> (Army, Army)
 170 handleOrder (Immune k) allocations _ attackArmy defendArmy = (attackArmy, defendArmy')
 171     where defendArmy' = handleAttack k allocations attackArmy defendArmy
 172 handleOrder (Infection k) _ allocations defendArmy attackArmy = (defendArmy', attackArmy)
 173     where defendArmy' = handleAttack k allocations attackArmy defendArmy
 174
 175 handleAttack :: Int -> [Allocation] -> Army -> Army -> Army
 176 handleAttack attacker allocations attackArmy defendArmy =
 177     if not $ null attackersAllocations
 178     then M.insert defender defendGroup' defendArmy
 179     else defendArmy
 180     where attackersAllocations = filter (\a -> attacker == fst a ) allocations
 181           defender = snd $ head attackersAllocations
 182           defendGroup = (defendArmy!defender)
 183           damage = damageCaused (attackArmy!attacker) defendGroup
 184           defendGroup' = applyDamage defendGroup damage
 185
 186
 187 battle :: Army -> Army -> (Army, Army)
 188 battle immuneArmy infectionArmy = head $ dropWhile (not . uncurry battleOver) $ iterate battleRound (immuneArmy, infectionArmy)
 189
 190
 191 pruneArmy :: Army -> Army
 192 pruneArmy = M.filter (\g -> _units g > 0)
 193
 194 battleOver :: Army -> Army -> Bool
 195 battleOver immuneArmy infectionArmy = (M.null immuneArmy) || (M.null infectionArmy)
 196
 197
 198 remainingUnitCount :: Army -> Army -> Int
 199 -- remainingUnitCount immuneArmy infectionArmy | trace ("End with\n" ++ show immuneArmy ++ " " ++ show infectionArmy) False = undefined
 200 remainingUnitCount immuneArmy infectionArmy = (unitCount immuneArmy) + (unitCount infectionArmy)
 201
 202 unitCount :: Army -> Int
 203 unitCount army = sum [_units g | g <- M.elems army]
 204
 205
 206 immuneWins :: (Army, Army) -> Bool
 207 immuneWins (immuneArmy, infectionArmy) = (unitCount immuneArmy > 0) && (unitCount infectionArmy == 0)
 208
 209
 210 applyBoost :: Int -> Army -> Army
 211 applyBoost boost = M.map (\g -> g { _damage = (_damage g + boost)})
 212
 213
 214 type Parser = Parsec Void Text
 215
 216 sc :: Parser ()
 217 sc = L.space (skipSome spaceChar) CA.empty CA.empty
 218
 219 lexeme  = L.lexeme sc
 220 integer = lexeme L.decimal
 221 -- signedInteger = L.signed sc integer
 222 symb = L.symbol sc
 223 comma = symb ","
 224 semicolon = symb ";"
 225 openBracket = symb "("
 226 closeBracket = symb ")"
 227
 228 immuneHeaderP = symb "Immune System:"
 229 infectionHeaderP = symb "Infection:"
 230
 231 sizePaddingP = symb "units each with"
 232 hpPaddingP = symb "hit points"
 233 attackPaddingP = symb "with an attack that does"
 234 initiativePaddingP = symb "damage at initiative"
 235 weaknessPaddingP = symb "weak to"
 236 immunityPaddingP = symb "immune to"
 237
 238
 239 armiesP = (,) <$> immuneGroupsP <*> infectionGroupsP
 240
 241 immuneGroupsP = immuneHeaderP *> many groupP
 242 infectionGroupsP = infectionHeaderP *> many groupP
 243
 244 -- 72 units each with 5294 hit points (weak to slashing; immune to radiation, cold) with an attack that does 639 fire damage at initiative 1
 245
 246 groupP = engroup <$> (integer <* sizePaddingP )
 247                  <*> (integer <* hpPaddingP)
 248                  <*> (attackModifierGroupP <* attackPaddingP )
 249                  <*> integer
 250                  <*> (damageTypeP <* initiativePaddingP)
 251                  <*> integer
 252     where engroup units hps aMods damage damageType initative =
 253                 Group { _units = units
 254                       , _hps = hps
 255                       , _modifiers = aMods
 256                       , _damage = damage
 257                       , _damageType = damageType
 258                       , _initiative = initative
 259                       }
 260
 261
 262 attackModifierGroupP = option [] ((openBracket `between` closeBracket) attackModifiersP)
 263
 264 attackModifiersP = attackModifierP `sepBy` semicolon
 265 attackModifierP = weaknessP <|> immunityP
 266 weaknessP = Weakness <$> (weaknessPaddingP *> damageTypesP)
 267 immunityP = Immunity <$> (immunityPaddingP *> damageTypesP)
 268
 269 damageTypeP = some letterChar <* sc
 270 damageTypesP = damageTypeP `sepBy` comma
 271
 272 successfulParse :: Text -> (Army, Army)
 273 -- successfulParse _ = []
 274 successfulParse input =
 275         case parse armiesP "input" input of
 276                 Left  _error -> (M.empty, M.empty) -- TIO.putStr $ T.pack $ parseErrorPretty err
 277                 Right armies -> idTag armies
 278         where idTag (immune, infect) = (tagArmy immune, tagArmy infect)
 279               tagArmy army = M.fromList $ zip [1..] army