module Main (main) where

import Data.List (transpose)

-- | Part 1
countXmas :: [Char] -> Int
countXmas [] = 0
countXmas s
  | take 4 s `elem` ["XMAS", "SAMX"] = 1 + countXmas (drop 1 s)
  | otherwise = countXmas (drop 1 s)

diagonals :: [[Char]] -> [[Char]]
diagonals grid =
  [[(grid !! (row + x)) !! (col + x) | x <- [0 .. 3]] | row <- [0 .. length grid - 4], col <- [0 .. length (grid !! row) - 4]]
    ++ [[(grid !! (row + x)) !! (col - x) | x <- [0 .. 3]] | row <- [0 .. length grid - 4], col <- [3 .. length (grid !! row) - 1]]

countAllXmas :: [[Char]] -> Int
countAllXmas grid = sum (map countXmas grid) + sum (map countXmas ((transpose . reverse) grid)) + sum (map countXmas (diagonals grid))

part1 :: String -> IO ()
part1 s = print $ countAllXmas $ lines s

-- | Part 2
subGrid3 :: [[Char]] -> Int -> Int -> [[Char]]
subGrid3 grid row col = [[(grid !! row') !! col' | col' <- [col .. col + 2]] | row' <- [row .. row + 2]]

masKernel :: [[Char]] -> Bool
masKernel [['M', _, 'M'], [_, 'A', _], ['S', _, 'S']] = True
masKernel [['M', _, 'S'], [_, 'A', _], ['M', _, 'S']] = True
masKernel [['S', _, 'M'], [_, 'A', _], ['S', _, 'M']] = True
masKernel [['S', _, 'S'], [_, 'A', _], ['M', _, 'M']] = True
masKernel _ = False

part2 :: String -> IO ()
part2 s =
  let grid = lines s
   in print $ length $ filter id [masKernel (subGrid3 grid row col) | row <- [0 .. length grid - 3], col <- [0 .. length (grid !! row) - 3]]

-- | Main
main :: IO ()
main = readFile "resources/day4.txt" >>= part2