{-# LANGUAGE LambdaCase #-}
module Main where

import Data.Functor
import Debug.Trace (trace)

data Cell where
  Paper :: Cell
  Empty :: Cell
    deriving (Eq, Show)

type Grid = [[Cell]]
type Pos = (Int,Int)
type GridWithPos = (Pos, [[Cell]])

joinCell :: Cell -> Cell -> Cell
joinCell Empty _     = Empty
joinCell _     Empty = Empty
joinCell Paper Paper = Paper

replace :: (Int,Int) -> a -> [[a]] -> [[a]]
replace (y,x) e es = take y es ++ replacex x e (es !!y):drop (y+1) es
  where replacex :: Int -> a -> [a] -> [a]
        replacex n e es = take n es ++ e:drop (n+1) es

parse :: String -> Grid
parse =
  (\x -> let n = length (head x)
             pad = replicate n Empty in
           pad : x <> pure pad) .
  map (\x -> Empty:x <> pure Empty) .
  map ((\case '@' -> Paper
              '.' -> Empty) <$>) . lines

solve1 :: Grid -> Int
solve1 = length . filter isAround . createWindows 3
  where
    createWindows :: Int -> Grid -> [Grid]
    createWindows n g =
      let h = length g
          w = length (head g)
          k = ((take n . ) . drop) . (+ (-1))
      in
        [1..h-n+1] >>= \y -> [1..w-n+1] >>= \x ->
        pure . map (k x) $ k y g

    isAround :: Grid -> Bool
    isAround g = if (g !! 1) !! 1 == Empty then False else
      (< (5 :: Int)) . sum $ g >>= \m -> m <&> \x -> if x == Empty then 0 else 1

solve2 :: Grid -> Int
solve2 = solve 0
  where
    solve :: Int -> Grid -> Int
    solve acc g =
      let newWindows = createWindows 3 g
          filtered = filter (isAround . snd) newWindows
          positions = map fst filtered
      in if null filtered then acc else
        solve
          (acc + length filtered)
          (foldr (`replace` Empty) g positions)

    createWindows :: Int -> Grid -> [GridWithPos]
    createWindows n g =
      let h = length g
          w = length (head g)
          k = ((take n . ) . drop) . (+ (-1))
      in
        [1..h-n+1] >>= \y -> [1..w-n+1] >>= \x ->
        pure . ((y,x),) . map (k x) $ k y g

    isAround :: Grid -> Bool
    isAround g = if (g !! 1) !! 1 == Empty then False else
      (< (5 :: Int)) . sum $ g >>= \m -> m <&> \x -> if x == Empty then 0 else 1


main :: IO ()
main = readFile "inputs/4" <&> parse >>= \i ->
       print (solve1 i) >>
       print (solve2 i)