import Aoc

parseLine :: String -> [Bool]
parseLine = map toBool
    where
        toBool '.' = False
        toBool '#' = True

countNeighbours :: [[Bool]] -> (Bool, Int)
countNeighbours [[a, b, c], [d, e, f], [g, h, i]] = (e, length (filter id [a, b, c, d, f, g, h, i]))

survives :: (Bool, Int) -> Bool
survives (False, n) = n == 3
survives (True, n) = n == 2 || n == 3

simulate :: [[Bool]] -> [[Bool]]
simulate = kernelMap (survives . countNeighbours) (Just False)

numOn :: [[Bool]] -> Int
numOn = sum . map (sum . map fromEnum)

doSimulate :: ([[Bool]] -> [[Bool]]) -> [[Bool]] -> Int
doSimulate simulate board = numOn $ iterate simulate board !! 100

part1 :: [[Bool]] -> Int
part1 = doSimulate simulate

part2 :: [[Bool]] -> Int
part2 = doSimulate simulate' . keepOn
    where
        updateFirst f (x:xs) = f x: xs
        updateLast f (x:xs) = if null xs then [f x] else x: updateLast f xs
        updateEnds f = updateFirst f . updateLast f
        keepOn = updateEnds (updateEnds $ const True)
        simulate' = keepOn . simulate

main :: IO ()
main = aocMain (map parseLine) part1 part2