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

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

Simulating a sea of octopuses that flash when they get excited.

-}
module Main (main) where

import Advent (count, getInputMap)
import Advent.Coord (Coord(..), neighbors)
import Data.Char (digitToInt)
import Data.List (elemIndex, foldl')
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Maybe (fromJust)

-- | >>> :main
-- 1585
-- 382
main :: IO ()
IO ()
main =
 do Map Coord Int
inp <- (Char -> Int) -> Map Coord Char -> 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 Char -> Int
digitToInt (Map Coord Char -> Map Coord Int)
-> IO (Map Coord Char) -> IO (Map Coord Int)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Int -> IO (Map Coord Char)
getInputMap Int
2021 Int
11
    let flashes :: [Int]
flashes = Map Coord Int -> [Int]
simulate Map Coord Int
inp
    Int -> IO ()
forall a. Show a => a -> IO ()
print ([Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum (Int -> [Int] -> [Int]
forall a. Int -> [a] -> [a]
take Int
100 [Int]
flashes))
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Maybe Int -> Int
forall a. HasCallStack => Maybe a -> a
fromJust (Int -> [Int] -> Maybe Int
forall a. Eq a => a -> [a] -> Maybe Int
elemIndex (Map Coord Int -> Int
forall k a. Map k a -> Int
Map.size Map Coord Int
inp) [Int]
flashes))

-- | Initial grid state to flashes per step
simulate :: Map Coord Int -> [Int]
simulate :: Map Coord Int -> [Int]
simulate = (Map Coord Int -> Int) -> [Map Coord Int] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map (Int -> Map Coord Int -> Int
forall (f :: * -> *) a. (Foldable f, Eq a) => a -> f a -> Int
count Int
0) ([Map Coord Int] -> [Int])
-> (Map Coord Int -> [Map Coord Int]) -> Map Coord Int -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Map Coord Int] -> [Map Coord Int]
forall a. HasCallStack => [a] -> [a]
tail ([Map Coord Int] -> [Map Coord Int])
-> (Map Coord Int -> [Map Coord Int])
-> Map Coord Int
-> [Map Coord Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Map Coord Int -> Map Coord Int)
-> Map Coord Int -> [Map Coord Int]
forall a. (a -> a) -> a -> [a]
iterate Map Coord Int -> Map Coord Int
step

-- | Advance the state of the world one time step
step :: Map Coord Int -> Map Coord Int
step :: Map Coord Int -> Map Coord Int
step Map Coord Int
m = (Map Coord Int -> Coord -> Map Coord Int)
-> Map Coord Int -> [Coord] -> Map Coord Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map Coord Int -> Coord -> Map Coord Int
excite ((Int -> Int) -> Map Coord Int -> 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 (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+) Map Coord Int
m) [Coord
k | (Coord
k, Int
9) <- Map Coord Int -> [(Coord, Int)]
forall k a. Map k a -> [(k, a)]
Map.toList Map Coord Int
m]

-- | Excite an octopus at the given location
excite :: Map Coord Int -> Coord -> Map Coord Int
excite :: Map Coord Int -> Coord -> Map Coord Int
excite Map Coord Int
m Coord
x =
  case Coord -> Map Coord Int -> Maybe Int
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Coord
x Map Coord Int
m of
    Just Int
e
      | Int
e Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
9 -> (Map Coord Int -> Coord -> Map Coord Int)
-> Map Coord Int -> [Coord] -> Map Coord Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Map Coord Int -> Coord -> Map Coord Int
excite (Coord -> Int -> Map Coord Int -> Map Coord Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Coord
x Int
0 Map Coord Int
m) (Coord -> [Coord]
neighbors Coord
x)
      | Int
e Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
1 -> Coord -> Int -> Map Coord Int -> Map Coord Int
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert Coord
x (Int
1 Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
e) Map Coord Int
m
    Maybe Int
_          -> Map Coord Int
m