module Main where

import Data.List (foldl')
import Data.Map (Map)
import qualified Data.Map as M
import System.Environment (getArgs)

main :: IO ()
main = do
  input <- readFile . head =<< getArgs
  putStrLn $ unwords ["Part 1:", show (part1 input)]
  putStrLn $ unwords ["Part 2:", show (part2 input)]

part1 :: String -> Int
part1 s = length $ filter (canPickRoll m (gridSize m)) (M.keys $ M.filter id m)
  where
    m = grid s

grid :: String -> Map (Int, Int) Bool
grid = M.fromList . concat . map (\(c, l) -> zipWith (\r ch -> ((r, c), ch == '@')) [0 ..] l) . zip [0 ..] . lines

-- | The size of a grid. Assumes that the grid is quadratical.
gridSize :: Map k a -> (Int, Int)
gridSize m = (\x -> (x, x)) . round . sqrt @Double . fromIntegral . M.size $ m

-- | A roll can be picked if less than 4 neighbors around it contain a
-- paper roll.
canPickRoll :: Map (Int, Int) Bool -> (Int, Int) -> (Int, Int) -> Bool
canPickRoll m s = (< 4) . length . filter (m M.!) . filter (validPos s) . neighbors
  where
    neighbors :: (Int, Int) -> [(Int, Int)]
    neighbors (r, c) =
      [ (r + 1, c + 1),
        (r + 1, c),
        (r + 1, c - 1),
        (r, c + 1),
        (r, c - 1),
        (r - 1, c + 1),
        (r - 1, c),
        (r - 1, c - 1)
      ]

    validPos :: (Int, Int) -> (Int, Int) -> Bool
    validPos (nr, nc) (r, c) = r >= 0 && r < nr && c >= 0 && c < nc

freeRolls :: (Int, Int) -> Map (Int, Int) Bool -> [(Int, Int)]
freeRolls s m = filter (canPickRoll m s) (M.keys $ M.filter id m)

-- | Remove all free rolls until there are no more free rolls and keep
-- a count of all the removed rolls.
part2 :: String -> Int
part2 input = removeFreeRolls (gridSize m) 0 m
  where
    m = grid input

removeFreeRolls :: (Int, Int) -> Int -> Map (Int, Int) Bool -> Int
removeFreeRolls s acc m = case freeRolls s m of
  [] -> acc
  free -> removeFreeRolls s (acc + length free) (foldl' (\m' k -> M.insert k False m') m free)

testInput :: String
testInput =
  unlines
    [ "..@@.@@@@.",
      "@@@.@.@.@@",
      "@@@@@.@.@@",
      "@.@@@@..@.",
      "@@.@@@@.@@",
      ".@@@@@@@.@",
      ".@.@.@.@@@",
      "@.@@@.@@@@",
      ".@@@@@@@@.",
      "@.@.@@@.@."
    ]