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

<https://adventofcode.com/2020/day/11>

-}
module Main (main) where

import Advent
import Advent.Coord
import Data.Array.Base qualified as AB
import Data.Array.Unboxed qualified as A
import Data.Maybe (mapMaybe)

type Seating   = A.UArray Coord Char

-- | Neighbors are stored using underlying raw indexes for fast access
-- on the 'Seating' grid.
type Neighbors = A.Array Coord [Int]

-- |
-- >>> :main
-- 2211
-- 1995
main :: IO ()
IO ()
main =
  do UArray Coord Char
inp <- Int -> Int -> IO (UArray Coord Char)
getInputArray Int
2020 Int
11
     let run :: (UArray Coord Char -> Maybe (UArray Coord Char)) -> IO ()
run UArray Coord Char -> Maybe (UArray Coord Char)
f = Int -> IO ()
forall a. Show a => a -> IO ()
print (Char -> [Char] -> Int
forall (f :: * -> *) a. (Foldable f, Eq a) => a -> f a -> Int
count Char
'#' (UArray Coord Char -> [Char]
forall (a :: * -> * -> *) e i. (IArray a e, Ix i) => a i e -> [e]
A.elems ((UArray Coord Char -> Maybe (UArray Coord Char))
-> UArray Coord Char -> UArray Coord Char
forall a. (a -> Maybe a) -> a -> a
stable UArray Coord Char -> Maybe (UArray Coord Char)
f UArray Coord Char
inp)))
     (UArray Coord Char -> Maybe (UArray Coord Char)) -> IO ()
run (Int -> Neighbors -> UArray Coord Char -> Maybe (UArray Coord Char)
adv Int
4 (UArray Coord Char -> Neighbors
adjacent UArray Coord Char
inp))
     (UArray Coord Char -> Maybe (UArray Coord Char)) -> IO ()
run (Int -> Neighbors -> UArray Coord Char -> Maybe (UArray Coord Char)
adv Int
5 (UArray Coord Char -> Neighbors
lineOfSight UArray Coord Char
inp))

-- | Repeatedly apply the function until it returns 'Nothing'. Return the
-- argument that returned 'Nothing'.
stable :: (a -> Maybe a) -> a -> a
stable :: forall a. (a -> Maybe a) -> a -> a
stable a -> Maybe a
f a
x = a -> (a -> a) -> Maybe a -> a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe a
x ((a -> Maybe a) -> a -> a
forall a. (a -> Maybe a) -> a -> a
stable a -> Maybe a
f) (a -> Maybe a
f a
x)

-- | Immediate neighbors used in part 1
adjacent :: Seating -> Neighbors
adjacent :: UArray Coord Char -> Neighbors
adjacent UArray Coord Char
a = (Coord, Coord) -> [[Int]] -> Neighbors
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
A.listArray (Coord, Coord)
b [[(Coord, Coord) -> Coord -> Int
forall a. Ix a => (a, a) -> a -> Int
A.index (Coord, Coord)
b Coord
j | Coord
j <- Coord -> [Coord]
neighbors Coord
i, (Coord, Coord) -> Coord -> Bool
forall a. Ix a => (a, a) -> a -> Bool
A.inRange (Coord, Coord)
b Coord
j] | Coord
i <- (Coord, Coord) -> [Coord]
forall a. Ix a => (a, a) -> [a]
A.range (Coord, Coord)
b]
  where
    b :: (Coord, Coord)
b = UArray Coord Char -> (Coord, Coord)
forall i. Ix i => UArray i Char -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
A.bounds UArray Coord Char
a

-- | Line of sight neighbors used in part 2
lineOfSight :: Seating -> Neighbors
lineOfSight :: UArray Coord Char -> Neighbors
lineOfSight UArray Coord Char
a = (Coord, Coord) -> [[Int]] -> Neighbors
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
(i, i) -> [e] -> a i e
A.listArray (Coord, Coord)
b [(Coord -> Maybe Int) -> [Coord] -> [Int]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (Coord -> Coord -> Maybe Int
look Coord
i) (Coord -> [Coord]
neighbors Coord
origin) | Coord
i <- (Coord, Coord) -> [Coord]
forall a. Ix a => (a, a) -> [a]
A.range (Coord, Coord)
b]
  where
    b :: (Coord, Coord)
b = UArray Coord Char -> (Coord, Coord)
forall i. Ix i => UArray i Char -> (i, i)
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> (i, i)
A.bounds UArray Coord Char
a
    look :: Coord -> Coord -> Maybe Int
look Coord
i Coord
d =
      do let j :: Coord
j = Coord
i Coord -> Coord -> Coord
forall a. Num a => a -> a -> a
+ Coord
d
         Char
v <- UArray Coord Char -> Coord -> Maybe Char
forall (a :: * -> * -> *) e i (f :: * -> *).
(IArray a e, Ix i, Alternative f) =>
a i e -> i -> f e
arrIx UArray Coord Char
a Coord
j
         case Char
v of
           Char
'.' -> Coord -> Coord -> Maybe Int
look Coord
j Coord
d
           Char
_   -> Int -> Maybe Int
forall a. a -> Maybe a
Just ((Coord, Coord) -> Coord -> Int
forall a. Ix a => (a, a) -> a -> Int
A.index (Coord, Coord)
b Coord
j)

-- | Advance the seating grid one time step using a configurable
-- threshold for seats becoming unoccupied, a precomputed neighborhood,
-- and the current seating chart. Return 'Nothing' when nothing changes.
adv ::
  Int           {- ^ occupied neighbor threshold      -} ->
  Neighbors     {- ^ neighborhood for each coordinate -} ->
  Seating       {- ^ current seating grid             -} ->
  Maybe Seating {- ^ updated seating grid             -}
adv :: Int -> Neighbors -> UArray Coord Char -> Maybe (UArray Coord Char)
adv Int
t Neighbors
ns UArray Coord Char
a
  | [(Coord, Char)] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(Coord, Char)]
changes = Maybe (UArray Coord Char)
forall a. Maybe a
Nothing
  | Bool
otherwise    = UArray Coord Char -> Maybe (UArray Coord Char)
forall a. a -> Maybe a
Just (UArray Coord Char -> Maybe (UArray Coord Char))
-> UArray Coord Char -> Maybe (UArray Coord Char)
forall a b. (a -> b) -> a -> b
$! UArray Coord Char
a UArray Coord Char -> [(Coord, Char)] -> UArray Coord Char
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)] -> a i e
A.// [(Coord, Char)]
changes
  where
    changes :: [(Coord, Char)]
changes = [(Coord
i, Char
v) | (Coord
i, Char
e) <- UArray Coord Char -> [(Coord, Char)]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> [(i, e)]
A.assocs UArray Coord Char
a, Char
v <- Coord -> Char -> [Char]
valueAt Coord
i Char
e]

    -- returns True when /at least/ n neighbors are occupied
    occupied :: Int -> Coord -> Bool
    occupied :: Int -> Coord -> Bool
occupied Int
n Coord
i = Int -> [Int] -> Bool
forall {t}. (Eq t, Num t) => t -> [Int] -> Bool
occupied1 Int
n (Neighbors
ns Neighbors -> Coord -> [Int]
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> i -> e
A.! Coord
i)

    occupied1 :: t -> [Int] -> Bool
occupied1 t
0 [Int]
_  = Bool
True
    occupied1 t
_ [] = Bool
False
    occupied1 t
n (Int
i:[Int]
is) =
      case UArray Coord Char -> Int -> Char
forall i. Ix i => UArray i Char -> Int -> Char
forall (a :: * -> * -> *) e i.
(IArray a e, Ix i) =>
a i e -> Int -> e
AB.unsafeAt UArray Coord Char
a Int
i of
        Char
'#' -> t -> [Int] -> Bool
occupied1 (t
nt -> t -> t
forall a. Num a => a -> a -> a
-t
1) [Int]
is
        Char
_   -> t -> [Int] -> Bool
occupied1 t
n [Int]
is

    valueAt :: Coord -> Char -> [Char]
valueAt Coord
i Char
'#' | Int -> Coord -> Bool
occupied Int
t Coord
i       = [Char]
"L"
    valueAt Coord
i Char
'L' | Bool -> Bool
not (Int -> Coord -> Bool
occupied Int
1 Coord
i) = [Char]
"#"
    valueAt Coord
_ Char
_                        = [Char]
""