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

<https://adventofcode.com/2022/day/23>

>>> :{
:main +
  ".......#......\n\
  \.....###.#....\n\
  \...#...#.#....\n\
  \....#...##....\n\
  \...#.###......\n\
  \...##.#.##....\n\
  \....#..#......\n"
:}
110
20

-}
module Main where

import Data.Array.Unboxed (Ix(rangeSize), UArray, accumArray)
import Data.List (foldl', tails)
import Data.Map qualified as Map
import Data.Maybe (fromMaybe)
import Data.Set (Set)
import Data.Set qualified as Set

import Advent (arrIx, getInputMap)
import Advent.Coord (Coord, above, below, boundingBox, left, neighbors, right)

-- |
-- >>> :main
-- 4236
-- 1023
main :: IO ()
IO ()
main =
 do Set Coord
elves <- Map Coord Char -> Set Coord
forall k a. Map k a -> Set k
Map.keysSet (Map Coord Char -> Set Coord)
-> (Map Coord Char -> Map Coord Char)
-> Map Coord Char
-> Set Coord
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> Map Coord Char -> Map Coord Char
forall a k. (a -> Bool) -> Map k a -> Map k a
Map.filter (Char
'#'Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
==) (Map Coord Char -> Set Coord)
-> IO (Map Coord Char) -> IO (Set Coord)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> IO (Map Coord Char)
getInputMap Int
2022 Int
23
    let states :: [Set Coord]
states = Set Coord -> [Set Coord]
sim Set Coord
elves

    -- part 1
    let round10 :: Set Coord
round10 = [Set Coord]
states [Set Coord] -> Int -> Set Coord
forall a. HasCallStack => [a] -> Int -> a
!! Int
10
    Int -> IO ()
forall a. Show a => a -> IO ()
print case Set Coord -> Maybe (Coord, Coord)
forall (f :: * -> *). Foldable f => f Coord -> Maybe (Coord, Coord)
boundingBox Set Coord
round10 of
      Just (Coord, Coord)
box -> (Coord, Coord) -> Int
forall a. Ix a => (a, a) -> Int
rangeSize (Coord, Coord)
box Int -> Int -> Int
forall a. Num a => a -> a -> a
- Set Coord -> Int
forall a. Set a -> Int
Set.size Set Coord
round10
      Maybe (Coord, Coord)
Nothing  -> Int
0

    -- part 2
    (Int -> IO ()) -> Maybe Int -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ Int -> IO ()
forall a. Show a => a -> IO ()
print ([Set Coord] -> Maybe Int
forall a. Ord a => [a] -> Maybe Int
sameIx [Set Coord]
states)

-- | Generate an infinite list of the rounds of the elf movement rules
sim :: Set Coord -> [Set Coord]
sim :: Set Coord -> [Set Coord]
sim Set Coord
start = (Set Coord
 -> (UArray Coord Bool -> Coord -> Maybe Coord) -> Set Coord)
-> Set Coord
-> [UArray Coord Bool -> Coord -> Maybe Coord]
-> [Set Coord]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl Set Coord
-> (UArray Coord Bool -> Coord -> Maybe Coord) -> Set Coord
step Set Coord
start [UArray Coord Bool -> Coord -> Maybe Coord]
moves

-- | Apply a single round of the elf movement rules
step ::
  Set Coord {- ^ initial elf locations -} ->
  (UArray Coord Bool -> Coord -> Maybe Coord) {- ^ proposal rule -} ->
  Set Coord {- ^ final elf locations -}
step :: Set Coord
-> (UArray Coord Bool -> Coord -> Maybe Coord) -> Set Coord
step Set Coord
elves UArray Coord Bool -> Coord -> Maybe Coord
m = Set Coord -> Map Coord Coord -> Set Coord
forall (f :: * -> *) a.
(Foldable f, Ord a) =>
Set a -> f a -> Set a
subtractFromSet Set Coord
elves Map Coord Coord
targets Set Coord -> Set Coord -> Set Coord
forall a. Semigroup a => a -> a -> a
<> Map Coord Coord -> Set Coord
forall k a. Map k a -> Set k
Map.keysSet Map Coord Coord
targets
   where
      !elves' :: UArray Coord Bool
elves' = Set Coord -> UArray Coord Bool
coordSet Set Coord
elves
      targets :: Map Coord Coord
targets = (Map Coord Coord -> Coord -> Map Coord Coord)
-> Map Coord Coord -> Set Coord -> Map Coord Coord
forall b a. (b -> a -> b) -> b -> Set a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map Coord Coord -> Coord -> Map Coord Coord
updateElf Map Coord Coord
forall k a. Map k a
Map.empty Set Coord
elves

      updateElf :: Map Coord Coord -> Coord -> Map Coord Coord
updateElf Map Coord Coord
acc Coord
src
        | UArray Coord Bool -> Coord -> Bool
isCrowded UArray Coord Bool
elves' Coord
src
        , Just Coord
dst <- UArray Coord Bool -> Coord -> Maybe Coord
m UArray Coord Bool
elves' Coord
src = (Maybe Coord -> Maybe Coord)
-> Coord -> Map Coord Coord -> Map Coord Coord
forall k a.
Ord k =>
(Maybe a -> Maybe a) -> k -> Map k a -> Map k a
Map.alter (Coord -> Maybe Coord -> Maybe Coord
forall {a} {a}. a -> Maybe a -> Maybe a
uniq Coord
src) Coord
dst Map Coord Coord
acc
        | Bool
otherwise = Map Coord Coord
acc

      uniq :: a -> Maybe a -> Maybe a
uniq a
v Maybe a
Nothing  = a -> Maybe a
forall a. a -> Maybe a
Just a
v  -- If the location is unassigned, assign it
      uniq a
_ (Just a
_) = Maybe a
forall a. Maybe a
Nothing -- If the location is assigned, unassign it

-- | Predicate testing to see if elf is near any other elf.
isCrowded :: CoordSet -> Coord -> Bool
isCrowded :: UArray Coord Bool -> Coord -> Bool
isCrowded UArray Coord Bool
elves Coord
elf = (Coord -> Bool) -> [Coord] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (UArray Coord Bool -> Coord -> Bool
arrayMember UArray Coord Bool
elves) (Coord -> [Coord]
neighbors Coord
elf)

-- | Move directions and their neighbors in the priority order of round 1
moveSets :: [(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
moveSets :: [(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
moveSets = [
  (Coord -> Coord
above, Coord -> Coord
left , Coord -> Coord
right),
  (Coord -> Coord
below, Coord -> Coord
left , Coord -> Coord
right),
  (Coord -> Coord
left , Coord -> Coord
above, Coord -> Coord
below),
  (Coord -> Coord
right, Coord -> Coord
above, Coord -> Coord
below)]

moves :: [CoordSet -> Coord -> Maybe Coord]
moves :: [UArray Coord Bool -> Coord -> Maybe Coord]
moves = ([(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
 -> UArray Coord Bool -> Coord -> Maybe Coord)
-> [[(Coord -> Coord, Coord -> Coord, Coord -> Coord)]]
-> [UArray Coord Bool -> Coord -> Maybe Coord]
forall a b. (a -> b) -> [a] -> [b]
map ([(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
-> UArray Coord Bool -> Coord -> Maybe Coord
forall {t}.
[(t -> Coord, Coord -> Coord, Coord -> Coord)]
-> UArray Coord Bool -> t -> Maybe Coord
combine ([(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
 -> UArray Coord Bool -> Coord -> Maybe Coord)
-> ([(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
    -> [(Coord -> Coord, Coord -> Coord, Coord -> Coord)])
-> [(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
-> UArray Coord Bool
-> Coord
-> Maybe Coord
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int
-> [(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
-> [(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
forall a. Int -> [a] -> [a]
take Int
4) ([(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
-> [[(Coord -> Coord, Coord -> Coord, Coord -> Coord)]]
forall a. [a] -> [[a]]
tails ([(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
-> [(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
forall a. HasCallStack => [a] -> [a]
cycle [(Coord -> Coord, Coord -> Coord, Coord -> Coord)]
moveSets))
  where
    combine :: [(t -> Coord, Coord -> Coord, Coord -> Coord)]
-> UArray Coord Bool -> t -> Maybe Coord
combine [] UArray Coord Bool
_ t
_ = Maybe Coord
forall a. Maybe a
Nothing
    combine ((t -> Coord
a,Coord -> Coord
b,Coord -> Coord
c):[(t -> Coord, Coord -> Coord, Coord -> Coord)]
xs) UArray Coord Bool
elves t
here
      | Bool -> Bool
not ((Coord -> Bool) -> [Coord] -> Bool
forall (t :: * -> *) a. Foldable t => (a -> Bool) -> t a -> Bool
any (UArray Coord Bool -> Coord -> Bool
arrayMember UArray Coord Bool
elves) [Coord]
locs) = Coord -> Maybe Coord
forall a. a -> Maybe a
Just Coord
here'
      | Bool
otherwise = [(t -> Coord, Coord -> Coord, Coord -> Coord)]
-> UArray Coord Bool -> t -> Maybe Coord
combine [(t -> Coord, Coord -> Coord, Coord -> Coord)]
xs UArray Coord Bool
elves t
here
      where
        here' :: Coord
here' = t -> Coord
a t
here
        locs :: [Coord]
locs = [Coord
here', Coord -> Coord
b Coord
here', Coord -> Coord
c Coord
here']

-- List utilites

-- | Find index in list where element is the same as the previous element.
sameIx :: Ord a => [a] -> Maybe Int
sameIx :: forall a. Ord a => [a] -> Maybe Int
sameIx = Int -> [a] -> Maybe Int
forall {a} {t}. (Eq a, Num t) => t -> [a] -> Maybe t
go Int
1
  where
    go :: t -> [a] -> Maybe t
go !t
i (a
x:a
y:[a]
z)
      | a
x a -> a -> Bool
forall a. Eq a => a -> a -> Bool
== a
y    = t -> Maybe t
forall a. a -> Maybe a
Just t
i
      | Bool
otherwise = t -> [a] -> Maybe t
go (t
it -> t -> t
forall a. Num a => a -> a -> a
+t
1) (a
ya -> [a] -> [a]
forall a. a -> [a] -> [a]
:[a]
z)
    go t
_ [a]
_ = Maybe t
forall a. Maybe a
Nothing

-- Set utilities

subtractFromSet :: (Foldable f, Ord a) => Set a -> f a -> Set a
subtractFromSet :: forall (f :: * -> *) a.
(Foldable f, Ord a) =>
Set a -> f a -> Set a
subtractFromSet = (Set a -> a -> Set a) -> Set a -> f a -> Set a
forall b a. (b -> a -> b) -> b -> f a -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' ((a -> Set a -> Set a) -> Set a -> a -> Set a
forall a b c. (a -> b -> c) -> b -> a -> c
flip a -> Set a -> Set a
forall a. Ord a => a -> Set a -> Set a
Set.delete)

-- | A set of coordinates represented with an array for fast membership lookups
type CoordSet = UArray Coord Bool

-- | Build an array-based representation of the set of coordinates. Because our elves are
-- densely packed, and we do lots of membership tests, this representation is more efficient.
coordSet :: Set Coord -> CoordSet
coordSet :: Set Coord -> UArray Coord Bool
coordSet Set Coord
s = (Bool -> Bool -> Bool)
-> Bool -> (Coord, Coord) -> [(Coord, Bool)] -> UArray Coord Bool
forall (a :: * -> * -> *) e i e'.
(IArray a e, Ix i) =>
(e -> e' -> e) -> e -> (i, i) -> [(i, e')] -> a i e
accumArray (\Bool
_old Bool
new -> Bool
new) Bool
False (Coord, Coord)
b [(Coord
c, Bool
True) | Coord
c <- Set Coord -> [Coord]
forall a. Set a -> [a]
Set.toList Set Coord
s]
  where
    b :: (Coord, Coord)
b = (Coord, Coord) -> Maybe (Coord, Coord) -> (Coord, Coord)
forall a. a -> Maybe a -> a
fromMaybe (Coord
0,Coord
0) (Set Coord -> Maybe (Coord, Coord)
forall (f :: * -> *). Foldable f => f Coord -> Maybe (Coord, Coord)
boundingBox Set Coord
s)

-- | Membership operation for CoordSet
arrayMember :: CoordSet -> Coord -> Bool
arrayMember :: UArray Coord Bool -> Coord -> Bool
arrayMember UArray Coord Bool
a Coord
x = Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
fromMaybe Bool
False (UArray Coord Bool -> Coord -> Maybe Bool
forall (a :: * -> * -> *) e i (f :: * -> *).
(IArray a e, Ix i, Alternative f) =>
a i e -> i -> f e
arrIx UArray Coord Bool
a Coord
x)