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

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

This problem has us implement a cellular automaton on an
infinite grid. The problem requires special treatment
so that we can represent updates to infinite space.

-}
module Main (main) where

import Advent (getInputLines, fromDigits, times)
import Advent.Coord(Coord(..), coordLines)
import Data.IntSet (IntSet)
import Data.IntSet qualified as IntSet
import Data.List (elemIndices)
import Data.Set (Set)
import Data.Set qualified as Set

-- | Pictures have the boolean value at all coordinates with the
-- exception of those listed in the set.
data Picture = Picture { Picture -> Bool
_background :: !Bool, Picture -> Set Coord
exceptions :: !(Set Coord) }

-- | >>> :main
-- 5179
-- 16112
main :: IO ()
main :: IO ()
main =
 do String
algStr : String
"" : [String]
imgStrs <- Int -> IO [String]
getInputLines Int
20
    let alg :: IntSet
alg = [Int] -> IntSet
IntSet.fromList (Char
'#' Char -> String -> [Int]
forall a. Eq a => a -> [a] -> [Int]
`elemIndices` String
algStr)
        pic :: Picture
pic = Bool -> Set Coord -> Picture
Picture Bool
False ([Coord] -> Set Coord
forall a. Ord a => [a] -> Set a
Set.fromList [Coord
c | (Coord
c, Char
'#') <- [String] -> [(Coord, Char)]
coordLines [String]
imgStrs])
        steps :: [Picture]
steps = (Picture -> Picture) -> Picture -> [Picture]
forall a. (a -> a) -> a -> [a]
iterate (IntSet -> Picture -> Picture
step IntSet
alg) Picture
pic
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Set Coord -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Picture -> Set Coord
exceptions ([Picture]
steps [Picture] -> Int -> Picture
forall a. [a] -> Int -> a
!!  Int
2)))
    Int -> IO ()
forall a. Show a => a -> IO ()
print (Set Coord -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length (Picture -> Set Coord
exceptions ([Picture]
steps [Picture] -> Int -> Picture
forall a. [a] -> Int -> a
!! Int
50)))

-- | Apply the given image enhancement algorithm to a picture
step :: IntSet -> Picture -> Picture
step :: IntSet -> Picture -> Picture
step IntSet
alg (Picture Bool
bg Set Coord
img) = Bool -> Set Coord -> Picture
Picture Bool
bg' Set Coord
img'
  where
    bg' :: Bool
bg' = Int -> IntSet -> Bool
IntSet.member (if Bool
bg then Int
0b111_111_111 else Int
0b000_000_000) IntSet
alg

    img' :: Set Coord
img' = (Coord -> Bool) -> Set Coord -> Set Coord
forall a. (a -> Bool) -> Set a -> Set a
Set.filter Coord -> Bool
except ([Coord] -> Set Coord
forall a. Ord a => [a] -> Set a
Set.fromList ((Coord -> [Coord]) -> Set Coord -> [Coord]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Coord -> [Coord]
cell Set Coord
img))

    except :: Coord -> Bool
except Coord
x = Int -> IntSet -> Bool
IntSet.member Int
n IntSet
alg Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
/= Bool
bg'
      where
          n :: Int
n = Int -> [Int] -> Int
forall a. Integral a => a -> [a] -> a
fromDigits Int
2 [if Coord -> Set Coord -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Coord
i Set Coord
img Bool -> Bool -> Bool
forall a. Eq a => a -> a -> Bool
/= Bool
bg then Int
1 else Int
0 | Coord
i <- Coord -> [Coord]
cell Coord
x]

-- | 3x3 neighborhood around a coordinate
cell :: Coord -> [Coord]
cell :: Coord -> [Coord]
cell (C Int
y Int
x) = Int -> Int -> Coord
C (Int -> Int -> Coord) -> [Int] -> [Int -> Coord]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> [Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1 .. Int
yInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1] [Int -> Coord] -> [Int] -> [Coord]
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> [Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
-Int
1 .. Int
xInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1]