{-# LANGUAGE LambdaCase, NamedFieldPuns, RecordWildCards, TypeApplications, ViewPatterns #-}
module Day15 (day15a, day15b) where
import Control.Arrow ((&&&), (***))
import Control.Parallel.Strategies (parBuffer, rseq, withStrategy)
import Data.Either (partitionEithers)
import Data.Foldable (asum)
import Data.Functor.Identity (runIdentity)
import Data.List (sortOn)
import qualified Data.List.NonEmpty as NE (fromList)
import Data.Map.Lazy (Map)
import qualified Data.Map.Lazy as M ((!?), alterF, elems, empty, filter, findWithDefault, fromDistinctAscList, fromListWith, insert, keys, keysSet, minViewWithKey, null, restrictKeys, toAscList, toList, union)
import Data.Maybe (fromJust)
import Data.Semigroup (Max(Max), Min(Min), sconcat)
import Data.Set (Set)
import qualified Data.Set as S ((\\), fromDistinctAscList, intersection, lookupMin, member, null, singleton, toList, union, unions)
import GHC.Conc (numCapabilities)
data Species = Elf | Goblin deriving (Eq)
data Unit e = Unit {species :: Species, hp :: e}
data Cave i e = Cave {walls :: Set (i, i), units :: Map (i, i) (Unit e)}
caveFromString :: (Integral i) => (Species -> e) -> String -> Cave i e
caveFromString initialHP string =
uncurry Cave . (S.fromDistinctAscList *** M.fromDistinctAscList) . partitionEithers $ do
(y, line) <- zip [0..] $ lines string
(x, c) <- zip [0..] line
case c of
'#' -> [Left (y, x)]
'E' -> [Right ((y, x), Unit Elf $ initialHP Elf)]
'G' -> [Right ((y, x), Unit Goblin $ initialHP Goblin)]
_ -> []
caveToDebugString :: (Integral i, Show e) => Cave i e -> String
caveToDebugString Cave {..} = unlines $ map rowToDebugString [y0..y1]
where ((Min y0, Max y1), (Min x0, Max x1)) = sconcat . NE.fromList $
((Min &&& Max) *** (Min &&& Max)) <$> S.toList walls
rowToDebugString y =
map (caveChar y) [x0..x1] ++ " " ++ M.findWithDefault "" y unitsByLine
caveChar y x
| S.member (y, x) walls = '#'
| Just Unit {species = Elf} <- units M.!? (y, x) = 'E'
| Just Unit {species = Goblin} <- units M.!? (y, x) = 'G'
| otherwise = '.'
unitsByLine = M.fromListWith (\a b -> b ++ ", " ++ a)
[ (y, (case species of Elf -> 'E'; Goblin -> 'G') : '(' : shows hp ")")
| ((y, _), Unit {..}) <- M.toAscList units
]
adjacencies :: (Num a) => (a, a) -> Set (a, a)
adjacencies (y, x) = S.fromDistinctAscList [(y - 1, x), (y, x - 1), (y, x + 1), (y + 1, x)]
nearest :: (Num a, Ord a) => Set (a, a) -> Set (a, a) -> (a, a) -> Maybe (a, a)
nearest walls goals = nearest' walls . S.singleton
where nearest' visited q
| S.null q = Nothing
| reached <- S.intersection goals q, not $ S.null reached = S.lookupMin reached
| otherwise = nearest' visited' $ S.unions (adjacencies <$> S.toList q) S.\\ visited
where visited' = S.union visited q
step :: (Monad m, Num i, Ord i, Ord e) => (Maybe (Unit e) -> m (Maybe (Unit e))) -> Set (i, i) -> Map (i, i) (Unit e) -> m (Map (i, i) (Unit e), Bool)
step strike walls = step' M.empty
where step' past (M.minViewWithKey -> Just ((k, unit), future))
| M.null enemies = return (M.insert k unit allOtherUnits, False)
| otherwise = case M.restrictKeys enemies $ adjacencies k' of
(sortOn (hp . snd) . M.toList -> (target, _):_) -> do
past' <- M.insert k' unit <$> M.alterF strike target past
future' <- M.alterF strike target future
step' past' future'
_ -> step' (M.insert k' unit past) future
where
allOtherUnits = M.union past future
walls' = S.union walls $ M.keysSet allOtherUnits
enemies = M.filter ((/= species unit) . species) allOtherUnits
adjacentEnemies = adjacencies k `S.intersection` M.keysSet enemies
enemyRanges = S.unions (adjacencies <$> M.keys enemies) S.\\ walls'
k' = case if S.null adjacentEnemies then nearest walls' enemyRanges k else Nothing of
Just goal | Just move <- nearest walls' (adjacencies k) goal -> move
_ -> k
step' past _ = return (past, True)
outcome :: (Monad m, Num i, Ord i, Num e, Ord e) => (Maybe (Unit e) -> m (Maybe (Unit e))) -> Cave i e -> m (Int, e)
outcome strike Cave {..} = outcome' 0 units
where outcome' rounds units' = step strike walls units' >>= \case
(units'', False) -> return (rounds, sum $ map hp $ M.elems units'')
(units'', True) -> outcome' (rounds + 1) units''
day15a :: String -> (Int, Int)
day15a (caveFromString @Int (const 200) -> cave) = runIdentity $ outcome strike cave
where strike (Just unit@Unit {hp}) | hp > 3 = return $ Just unit {hp = hp - 3}
strike _ = return Nothing
day15b :: String -> (Int, (Int, Int))
day15b (caveFromString @Int (const 200) -> cave) = fromJust $ asum $ parallelize
[(,) elfPower <$> outcome (strike elfPower) cave | elfPower <- [3..]]
where strike _ (Just unit@Unit {species = Elf, hp})
| hp > 3 = return $ Just unit {hp = hp - 3}
| otherwise = fail "oh noes!"
strike elfPower (Just unit@Unit {hp})
| hp > elfPower = return $ Just unit {hp = hp - elfPower}
strike _ _ = return Nothing
parallelize = withStrategy $ parBuffer numCapabilities rseq