module Day11 where
import Data.Array
    ( Ix, Array, (!), (//), assocs, bounds, listArray )
import Data.Char (digitToInt)
import Data.Maybe ( catMaybes, mapMaybe )
import Data.List (unfoldr)
-- import GHC.Ix (Ix)

data Point = P Int Int deriving (Eq, Ord, Ix, Show)

flashOctos:: [Point] -> Array Point Int -> [Point] -> (Array Point Int, [Point])
flashOctos [] octos points = (octos, points)
flashOctos (p@(P x y):xs) octos points = if p `elem` points || octos ! p <= 9
    then flashOctos xs octos points
    else flashOctos (xs ++ filter (\p -> octos ! p >= 9) neighs) (octos // increments) (p:points)
    where
        neighs = catMaybes $ (\x' y' -> if (x == x' && y == y') || x' < 1 || x' > mx || y' < 1 || y' > my then Nothing else Just (P x' y')) <$> [x-1..x+1] <*> [y-1..y+1]
        (P 1 1, P mx my) = bounds octos
        increments = map(\p -> (p, (octos ! p) + 1)) neighs

stepOctos :: Array Point Int -> (Array Point Int, Int)
stepOctos octos = (field // map (\p -> (p,0)) toReset, length toReset)
    where
        (field,toReset) = flashOctos points octos' []
        points = mapMaybe (\(p,x) -> if x > 9 then Just p else Nothing) $ assocs octos'
        octos' = fmap (+ 1) octos

part2 :: Array Point Int -> Int -> Int
part2 a i = if c == length a' then i else part2 a' (i+1)
    where (a', c) = stepOctos a

main :: IO ()
main = do
    input <- map (map digitToInt) . lines <$> getContents

    putStr "part 1: "
    let octos = listArray (P 1 1, P (length input) (length (head input))) (concat input)
    let (_,count) = foldr (\_ (o, count) -> let (o', c') = stepOctos o in (o', count+c')) (octos, 0) [1..100]
    print count

    putStr "part 2: " >> print (part2 octos 1)


    return ()