module Day20 where
import qualified Data.Array as A
import qualified Data.Set as S
import qualified Data.Map as M
import Debug.Trace (traceShow, trace)
import Data.Maybe (catMaybes)

type Algorithm = A.Array Int Bool
data Image = Img Bool [(Int,Int)] deriving (Show)

getNd :: Image -> [(Int,Int)]
getNd (Img _ s) = s

getLit :: Image -> Int
getLit (Img True s) = error "all elems lit"
getLit (Img False s) = length s

newDefault :: Algorithm -> Bool -> Bool
newDefault algo False = algo A.! 0
newDefault algo True  = algo A.! 511

insertManyWith :: Ord k => (a -> a -> a) -> [(k,a)] -> M.Map k a -> M.Map k a
insertManyWith f [] = id
insertManyWith f ((k,v):xs) = insertManyWith f xs . M.insertWith f k v

binToDec :: [Bool] -> Int
binToDec = foldr (\x acc -> acc * 2 + fromEnum x) 0

next :: Algorithm -> Image -> Image
next algo (Img def diff) = Img (newDefault algo def) newDiff where
    nonDefault = foldr (\(x,y) -> insertManyWith (++) (neighs x y)) M.empty diff
    neighs x y = map (\i -> ((x-1 + i `mod` 3, y-1 + i `div` 3), [i])) [0..8]
    known = M.map (\ns -> algo A.! binToDec (nonDefToNum ns)) nonDefault
    calcNonDef = M.filter (/= newDefault algo def) known
    nonDefToNum ns = map (\i -> if i `elem` ns then not def else def) [0..8]
    newDiff = M.keys calcNonDef

printImage :: Image -> IO ()
printImage img = do
    mapM_ (\y -> do
        mapM_ (\x -> if (x,y) `elem` getNd img then putChar '#' else putChar '.') [minx..maxx];
        putStrLn "") [miny..maxy] where
    minx = minimum $ map fst $ getNd img
    maxx = maximum $ map fst $ getNd img
    miny = minimum $ map snd $ getNd img
    maxy = maximum $ map snd $ getNd img


main :: IO()
main = do
    algorithm <- A.listArray (0,511) . map (=='#') <$> getLine
    getLine
    content <- lines <$> getContents
    let nonDef = catMaybes $ concat $ zipWith (\j -> zipWith (\i e -> if e =='#' then Just (i,j) else Nothing) [0..]) [0..] content
    let init = Img False nonDef


    let enhanced = iterate (next algorithm) init
    putStr "part 1: "
    print $ getLit $ enhanced !! 2

    -- let enhanced = iterate (next algorithm) init !! 50
    putStr "part 2: "
    print $ getLit $ enhanced !! 50



    print "bye"