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

<https://adventofcode.com/2018/day/15>
-}
module Main (main) where

import Advent ( getInputArray, same )
import Advent.Coord ( cardinal, Coord )
import Data.Array.Unboxed ( UArray, (!), amap, assocs )
import Data.List ( foldl', minimumBy )
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe ( fromMaybe )
import Data.Ord ( comparing )
import Data.Set (Set)
import Data.Set qualified as Set

data Team = Elf | Goblin
  deriving (Team -> Team -> Bool
(Team -> Team -> Bool) -> (Team -> Team -> Bool) -> Eq Team
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Team -> Team -> Bool
== :: Team -> Team -> Bool
$c/= :: Team -> Team -> Bool
/= :: Team -> Team -> Bool
Eq, Eq Team
Eq Team =>
(Team -> Team -> Ordering)
-> (Team -> Team -> Bool)
-> (Team -> Team -> Bool)
-> (Team -> Team -> Bool)
-> (Team -> Team -> Bool)
-> (Team -> Team -> Team)
-> (Team -> Team -> Team)
-> Ord Team
Team -> Team -> Bool
Team -> Team -> Ordering
Team -> Team -> Team
forall a.
Eq a =>
(a -> a -> Ordering)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> Bool)
-> (a -> a -> a)
-> (a -> a -> a)
-> Ord a
$ccompare :: Team -> Team -> Ordering
compare :: Team -> Team -> Ordering
$c< :: Team -> Team -> Bool
< :: Team -> Team -> Bool
$c<= :: Team -> Team -> Bool
<= :: Team -> Team -> Bool
$c> :: Team -> Team -> Bool
> :: Team -> Team -> Bool
$c>= :: Team -> Team -> Bool
>= :: Team -> Team -> Bool
$cmax :: Team -> Team -> Team
max :: Team -> Team -> Team
$cmin :: Team -> Team -> Team
min :: Team -> Team -> Team
Ord, Int -> Team -> ShowS
[Team] -> ShowS
Team -> String
(Int -> Team -> ShowS)
-> (Team -> String) -> ([Team] -> ShowS) -> Show Team
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Team -> ShowS
showsPrec :: Int -> Team -> ShowS
$cshow :: Team -> String
show :: Team -> String
$cshowList :: [Team] -> ShowS
showList :: [Team] -> ShowS
Show)

data Unit = Unit { Unit -> Int
attack, Unit -> Int
hp :: !Int, Unit -> Team
team :: !Team }
  deriving (Int -> Unit -> ShowS
[Unit] -> ShowS
Unit -> String
(Int -> Unit -> ShowS)
-> (Unit -> String) -> ([Unit] -> ShowS) -> Show Unit
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Unit -> ShowS
showsPrec :: Int -> Unit -> ShowS
$cshow :: Unit -> String
show :: Unit -> String
$cshowList :: [Unit] -> ShowS
showList :: [Unit] -> ShowS
Show)

type Dungeon = UArray Coord Bool

-- | Print the answers to day 15
--
-- >>> :main
-- 346574
-- 60864
main :: IO ()
IO ()
main =
  do UArray Coord Char
input <- Int -> Int -> IO (UArray Coord Char)
getInputArray Int
2018 Int
15
     let dungeon :: Dungeon
dungeon = UArray Coord Char -> Dungeon
parseMap UArray Coord Char
input
         units :: Map Coord Team
units = UArray Coord Char -> Map Coord Team
parseUnits UArray Coord Char
input

     Int -> IO ()
forall a. Show a => a -> IO ()
print (Dungeon -> Map Coord Team -> Int
part1 Dungeon
dungeon Map Coord Team
units)
     Int -> IO ()
forall a. Show a => a -> IO ()
print (Dungeon -> Map Coord Team -> Int
part2 Dungeon
dungeon Map Coord Team
units)

part1 :: Dungeon -> Map Coord Team -> Int
part1 :: Dungeon -> Map Coord Team -> Int
part1 Dungeon
dungeon Map Coord Team
units = Map Coord Unit -> Int -> Int
outcome Map Coord Unit
units2 Int
turns
  where
    units1 :: Map Coord Unit
units1 = (Team -> Unit) -> Map Coord Team -> Map Coord Unit
forall a b. (a -> b) -> Map Coord a -> Map Coord b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Int -> Int -> Team -> Unit
Unit Int
3 Int
200) Map Coord Team
units
    (Map Coord Unit
units2, Int
turns) = [(Map Coord Unit, Int)] -> (Map Coord Unit, Int)
forall a. HasCallStack => [a] -> a
last (Dungeon -> Map Coord Unit -> Int -> [(Map Coord Unit, Int)]
simulate Dungeon
dungeon Map Coord Unit
units1 Int
0)

part2 :: Dungeon -> Map Coord Team -> Int
part2 :: Dungeon -> Map Coord Team -> Int
part2 Dungeon
dungeon Map Coord Team
units = Int -> Int -> Maybe (Int, Int) -> Int
search Int
elfCount Int
3 Maybe (Int, Int)
forall a. Maybe a
Nothing
  where
    elfCount :: Int
elfCount = Map Coord Team -> Int
forall k a. Map k a -> Int
Map.size ((Team -> Bool) -> Map Coord Team -> Map Coord Team
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (Team
ElfTeam -> Team -> Bool
forall a. Eq a => a -> a -> Bool
==) Map Coord Team
units)

    search :: Int -> Int -> Maybe (Int, Int) -> Int
search Int
elfCount Int
atkLo (Just (Int
atkHi, Int
answer))
      | Int
atkLo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
1 Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
atkHi = Int
answer

    search Int
elfCount Int
atkLo Maybe (Int, Int)
mbAtkHi
      | [(Map Coord Unit, Int)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Map Coord Unit, Int)]
trimmed -- perfect elf victory
      = Int -> Int -> Maybe (Int, Int) -> Int
search Int
elfCount Int
atkLo ((Int, Int) -> Maybe (Int, Int)
forall a. a -> Maybe a
Just (Int
atk, Map Coord Unit -> Int -> Int
outcome Map Coord Unit
units2 Int
turns))
      | Bool
otherwise = Int -> Int -> Maybe (Int, Int) -> Int
search Int
elfCount Int
atk Maybe (Int, Int)
mbAtkHi
      where
        atk :: Int
atk = case Maybe (Int, Int)
mbAtkHi of
                Maybe (Int, Int)
Nothing       -> Int
atkLo Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
2
                Just (Int
atkHi, Int
_) -> (Int
atkLo Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
atkHi) Int -> Int -> Int
forall a. Integral a => a -> a -> a
`quot` Int
2
        toUnit :: Team -> Unit
toUnit Team
t = case Team
t of Team
Elf -> Int -> Int -> Team -> Unit
Unit Int
atk Int
200 Team
t; Team
Goblin -> Int -> Int -> Team -> Unit
Unit Int
3 Int
200 Team
t
        units1 :: Map Coord Unit
units1 = (Team -> Unit) -> Map Coord Team -> Map Coord Unit
forall a b. (a -> b) -> Map Coord a -> Map Coord b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Team -> Unit
toUnit Map Coord Team
units
        allElves :: (Map k Unit, b) -> Bool
allElves (Map k Unit
us,b
_) = Map k Unit -> Int
forall k a. Map k a -> Int
Map.size ((Unit -> Bool) -> Map k Unit -> Map k Unit
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (\Unit
u -> Unit -> Team
team Unit
u Team -> Team -> Bool
forall a. Eq a => a -> a -> Bool
== Team
Elf) Map k Unit
us)
                          Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
elfCount
        ([(Map Coord Unit, Int)]
ticks, [(Map Coord Unit, Int)]
trimmed) = ((Map Coord Unit, Int) -> Bool)
-> [(Map Coord Unit, Int)]
-> ([(Map Coord Unit, Int)], [(Map Coord Unit, Int)])
forall a. (a -> Bool) -> [a] -> ([a], [a])
span (Map Coord Unit, Int) -> Bool
forall {k} {b}. (Map k Unit, b) -> Bool
allElves (Dungeon -> Map Coord Unit -> Int -> [(Map Coord Unit, Int)]
simulate Dungeon
dungeon Map Coord Unit
units1 Int
0)

        (Map Coord Unit
units2, Int
turns) = [(Map Coord Unit, Int)] -> (Map Coord Unit, Int)
forall a. HasCallStack => [a] -> a
last [(Map Coord Unit, Int)]
ticks


outcome :: Map Coord Unit -> Int -> Int
outcome :: Map Coord Unit -> Int -> Int
outcome Map Coord Unit
units Int
turns = Int
turns Int -> Int -> Int
forall a. Num a => a -> a -> a
* Map Coord Int -> Int
forall a. Num a => Map Coord a -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Unit -> Int) -> Map Coord Unit -> Map Coord Int
forall a b. (a -> b) -> Map Coord a -> Map Coord b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Unit -> Int
hp Map Coord Unit
units)

parseMap :: UArray Coord Char -> Dungeon
parseMap :: UArray Coord Char -> Dungeon
parseMap = (Char -> Bool) -> UArray Coord Char -> Dungeon
forall (a :: * -> * -> *) e' e i.
(IArray a e', IArray a e, Ix i) =>
(e' -> e) -> a i e' -> a i e
amap (Char
'#' Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
/=)

parseUnits :: UArray Coord Char -> Map Coord Team
parseUnits :: UArray Coord Char -> Map Coord Team
parseUnits UArray Coord Char
rs = [(Coord, Team)] -> Map Coord Team
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList
  [ (Coord
c, Team
unit)
    | (Coord
c,Char
v) <- UArray Coord Char -> [(Coord, Char)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
assocs UArray Coord Char
rs
    , Team
unit <- case Char
v of
                Char
'G' -> [Team
Goblin]
                Char
'E' -> [Team
Elf   ]
                Char
_   -> [      ]
    ]

simulate :: Dungeon -> Map Coord Unit -> Int -> [(Map Coord Unit, Int)]
simulate :: Dungeon -> Map Coord Unit -> Int -> [(Map Coord Unit, Int)]
simulate Dungeon
dungeon Map Coord Unit
units Int
turns =
  (Map Coord Unit
units, Int
turns) (Map Coord Unit, Int)
-> [(Map Coord Unit, Int)] -> [(Map Coord Unit, Int)]
forall a. a -> [a] -> [a]
: Dungeon
-> Set Coord -> Map Coord Unit -> Int -> [(Map Coord Unit, Int)]
tick Dungeon
dungeon (Map Coord Unit -> Set Coord
forall k a. Map k a -> Set k
Map.keysSet Map Coord Unit
units) Map Coord Unit
units Int
turns

tick :: Dungeon -> Set Coord -> Map Coord Unit -> Int -> [(Map Coord Unit, Int)]
tick :: Dungeon
-> Set Coord -> Map Coord Unit -> Int -> [(Map Coord Unit, Int)]
tick Dungeon
dungeon Set Coord
schedule Map Coord Unit
units Int
turns =
  case Set Coord -> Maybe (Coord, Set Coord)
forall a. Set a -> Maybe (a, Set a)
Set.minView Set Coord
schedule of
    Maybe (Coord, Set Coord)
Nothing -> Dungeon -> Map Coord Unit -> Int -> [(Map Coord Unit, Int)]
simulate Dungeon
dungeon Map Coord Unit
units (Int
turnsInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
    Maybe (Coord, Set Coord)
_ | Map Coord Team -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => t a -> Bool
same ((Unit -> Team) -> Map Coord Unit -> Map Coord Team
forall a b. (a -> b) -> Map Coord a -> Map Coord b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Unit -> Team
team Map Coord Unit
units) -> [(Map Coord Unit
units,Int
turns)]
    Just (Coord
pos, Set Coord
schedule) ->
      let unit :: Unit
unit = Map Coord Unit
units Map Coord Unit -> Coord -> Unit
forall k a. Ord k => Map k a -> k -> a
Map.! Coord
pos
          pos' :: Coord
pos' = Coord -> Maybe Coord -> Coord
forall a. a -> Maybe a -> a
fromMaybe Coord
pos (Coord -> Unit -> Map Coord Unit -> Dungeon -> Maybe Coord
route Coord
pos Unit
unit (Coord -> Map Coord Unit -> Map Coord Unit
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Coord
pos Map Coord Unit
units) Dungeon
dungeon)
          units' :: Map Coord Unit
units' = Coord -> Unit -> Map Coord Unit -> Map Coord Unit
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Coord
pos' Unit
unit (Coord -> Map Coord Unit -> Map Coord Unit
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete Coord
pos Map Coord Unit
units)
      in case Coord -> Unit -> Map Coord Unit -> Maybe Coord
target Coord
pos' Unit
unit Map Coord Unit
units of
            Just Coord
tgt ->
               let units'' :: Map Coord Unit
units'' = Int -> Coord -> Map Coord Unit -> Map Coord Unit
melee (Unit -> Int
attack Unit
unit) Coord
tgt Map Coord Unit
units'
                   schedule' :: Set Coord
schedule' = Set Coord -> Set Coord -> Set Coord
forall a. Ord a => Set a -> Set a -> Set a
Set.intersection Set Coord
schedule (Map Coord Unit -> Set Coord
forall k a. Map k a -> Set k
Map.keysSet Map Coord Unit
units'')
               in Dungeon
-> Set Coord -> Map Coord Unit -> Int -> [(Map Coord Unit, Int)]
tick Dungeon
dungeon Set Coord
schedule' Map Coord Unit
units'' Int
turns
            Maybe Coord
Nothing ->
               Dungeon
-> Set Coord -> Map Coord Unit -> Int -> [(Map Coord Unit, Int)]
tick Dungeon
dungeon Set Coord
schedule Map Coord Unit
units' Int
turns

-- | Attack the unit at a coordinate for a given amount of damage.
-- Once a unit no longer has positive hit points it is removed from
-- the map of units.
melee ::
  Int            {- ^ damage              -} ->
  Coord          {- ^ target's coordinate -} ->
  Map Coord Unit {- ^ all units           -} ->
  Map Coord Unit {- ^ updated units       -}
melee :: Int -> Coord -> Map Coord Unit -> Map Coord Unit
melee Int
atk = (Unit -> Maybe Unit) -> Coord -> Map Coord Unit -> Map Coord Unit
forall k a. Ord k => (a -> Maybe a) -> k -> Map k a -> Map k a
Map.update ((Unit -> Maybe Unit) -> Coord -> Map Coord Unit -> Map Coord Unit)
-> (Unit -> Maybe Unit)
-> Coord
-> Map Coord Unit
-> Map Coord Unit
forall a b. (a -> b) -> a -> b
$ \Unit
tgt ->
  if Unit -> Int
hp Unit
tgt Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
atk
    then Maybe Unit
forall a. Maybe a
Nothing
    else Unit -> Maybe Unit
forall a. a -> Maybe a
Just (Unit -> Maybe Unit) -> Unit -> Maybe Unit
forall a b. (a -> b) -> a -> b
$! Unit
tgt { hp = hp tgt - atk }

-- | Figure out what neighboring unit this unit wants to attack
target ::
  Coord {- ^ coordinate of attacking unit -} ->
  Unit -> Map Coord Unit -> Maybe Coord
target :: Coord -> Unit -> Map Coord Unit -> Maybe Coord
target Coord
pos Unit
unit Map Coord Unit
units
  | [Coord] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Coord]
possible = Maybe Coord
forall a. Maybe a
Nothing
  | Bool
otherwise     = Coord -> Maybe Coord
forall a. a -> Maybe a
Just (Coord -> Maybe Coord) -> Coord -> Maybe Coord
forall a b. (a -> b) -> a -> b
$! (Coord -> Coord -> Ordering) -> [Coord] -> Coord
forall (t :: * -> *) a.
Foldable t =>
(a -> a -> Ordering) -> t a -> a
minimumBy Coord -> Coord -> Ordering
ordering [Coord]
possible
  where
    ordering :: Coord -> Coord -> Ordering
ordering = (Coord -> (Int, Coord)) -> Coord -> Coord -> Ordering
forall a b. Ord a => (b -> a) -> b -> b -> Ordering
comparing (\Coord
i -> (Unit -> Int
hp (Map Coord Unit
units Map Coord Unit -> Coord -> Unit
forall k a. Ord k => Map k a -> k -> a
Map.! Coord
i), Coord
i))
    possible :: [Coord]
possible = (Coord -> Bool) -> [Coord] -> [Coord]
forall a. (a -> Bool) -> [a] -> [a]
filter (Unit -> Map Coord Unit -> Coord -> Bool
isEnemy Unit
unit Map Coord Unit
units) (Coord -> [Coord]
cardinal Coord
pos)

-- | Determine if a given coordinate contains an enemy of the given unit.
isEnemy ::
  Unit           {- ^ focused unit            -} ->
  Map Coord Unit {- ^ all units               -} ->
  Coord          {- ^ possible enemy location -} ->
  Bool           {- ^ enemy at location       -}
isEnemy :: Unit -> Map Coord Unit -> Coord -> Bool
isEnemy Unit
unit Map Coord Unit
units Coord
loc =
  case Coord -> Map Coord Unit -> Maybe Unit
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Coord
loc Map Coord Unit
units of
    Just Unit
u  -> Unit -> Team
team Unit
u Team -> Team -> Bool
forall a. Eq a => a -> a -> Bool
/= Unit -> Team
team Unit
unit
    Maybe Unit
Nothing -> Bool
False

-- | Figure out where, if anywhere, this unit wants to move
route ::
  Coord          {- ^ unit's position -} ->
  Unit           {- ^ unit stats      -} ->
  Map Coord Unit {- ^ all units       -} ->
  Dungeon        {- ^ dungeon map     -} ->
  Maybe Coord    {- ^ next location   -}
route :: Coord -> Unit -> Map Coord Unit -> Dungeon -> Maybe Coord
route Coord
pos Unit
unit Map Coord Unit
units Dungeon
dungeon
  | Coord -> Bool
isNear Coord
pos = Maybe Coord
forall a. Maybe a
Nothing
  | Bool
otherwise  = Set Coord -> Set (Integer, Coord, Coord) -> Maybe Coord
forall {a} {b}.
(Ord a, Ord b, Num a) =>
Set Coord -> Set (a, Coord, b) -> Maybe b
search Set Coord
forall a. Set a
Set.empty Set (Integer, Coord, Coord)
candidates
  where
    candidates :: Set (Integer, Coord, Coord)
candidates = [(Integer, Coord, Coord)] -> Set (Integer, Coord, Coord)
forall a. Ord a => [a] -> Set a
Set.fromList [ (Integer
0, Coord
start, Coord
start) | Coord
start <- Coord -> [Coord]
cardinal Coord
pos ]
    isOpen :: Coord -> Bool
isOpen Coord
loc = Dungeon
dungeon Dungeon -> Coord -> Bool
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
! Coord
loc Bool -> Bool -> Bool
&& Coord -> Map Coord Unit -> Bool
forall k a. Ord k => k -> Map k a -> Bool
Map.notMember Coord
loc Map Coord Unit
units
    isNear :: Coord -> Bool
isNear = (Coord -> Bool) -> [Coord] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (Unit -> Map Coord Unit -> Coord -> Bool
isEnemy Unit
unit Map Coord Unit
units) ([Coord] -> Bool) -> (Coord -> [Coord]) -> Coord -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Coord -> [Coord]
cardinal

    search :: Set Coord -> Set (a, Coord, b) -> Maybe b
search Set Coord
seen Set (a, Coord, b)
q =
      do ((a
dist, Coord
dest, b
start), Set (a, Coord, b)
q) <- Set (a, Coord, b) -> Maybe ((a, Coord, b), Set (a, Coord, b))
forall a. Set a -> Maybe (a, Set a)
Set.minView Set (a, Coord, b)
q
         if | Coord -> Set Coord -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Coord
dest Set Coord
seen Bool -> Bool -> Bool
|| Bool -> Bool
not (Coord -> Bool
isOpen Coord
dest) -> Set Coord -> Set (a, Coord, b) -> Maybe b
search Set Coord
seen Set (a, Coord, b)
q
            | Coord -> Bool
isNear Coord
dest -> b -> Maybe b
forall a. a -> Maybe a
Just b
start
            | Bool
otherwise -> Set Coord -> Set (a, Coord, b) -> Maybe b
search (Coord -> Set Coord -> Set Coord
forall a. Ord a => a -> Set a -> Set a
Set.insert Coord
dest Set Coord
seen)
                                 ((Set (a, Coord, b) -> (a, Coord, b) -> Set (a, Coord, b))
-> Set (a, Coord, b) -> [(a, Coord, b)] -> Set (a, Coord, b)
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (((a, Coord, b) -> Set (a, Coord, b) -> Set (a, Coord, b))
-> Set (a, Coord, b) -> (a, Coord, b) -> Set (a, Coord, b)
forall a b c. (a -> b -> c) -> b -> a -> c
flip (a, Coord, b) -> Set (a, Coord, b) -> Set (a, Coord, b)
forall a. Ord a => a -> Set a -> Set a
Set.insert) Set (a, Coord, b)
q
                                     [ (a
dista -> a -> a
forall a. Num a => a -> a -> a
+a
1, Coord
dest, b
start) | Coord
dest <- Coord -> [Coord]
cardinal Coord
dest ])