module Main where

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

main :: IO ()
main = do
  input <- readFile . head =<< getArgs
  print $ part1 input
  print $ part2 input

part1 :: String -> Int
part1 s = length $ filter (canPickRoll m) (M.keys 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

canPickRoll :: Map (Int, Int) Bool -> (Int, Int) -> Bool
canPickRoll m p
  | m M.! p = checkNeighbor m p
  | otherwise = False
  where
    checkNeighbor :: Map (Int, Int) Bool -> (Int, Int) -> Bool
    checkNeighbor m' = (< 4) . length . filter (m' M.!) . filter (validPos gridSize) . neighbors

    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

    gridSize :: (Int, Int)
    gridSize = duplicate . round . sqrt @Double . fromIntegral . M.size $ m

    duplicate :: a -> (a, a)
    duplicate a = (a, a)

part2 :: String -> Int
part2 = error "Not implemented"

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