import Aoc
import Data.List.Split ( splitOn )
import Data.List ( transpose )

horizontal :: [String] -> Int
horizontal = sum . map horizontal'
    where 
        count word = pred . length . splitOn word
        horizontal' line = count "XMAS" line + count "SAMX" line

sheer :: [String] -> [String]
sheer [] = []
sheer (x:xs) = x:(map ('.':) $ sheer xs)

part1 :: [String] -> Int
part1 board = horizontal board + vertical board + diagonal1 board + diagonal2 board
    where
        vertical = horizontal . transpose
        diagonal1 = vertical . sheer
        diagonal2 = diagonal1 . reverse

isXMAS :: [[Char]] -> Bool
isXMAS [
    [a,  _,  b], 
    [_, 'A', _], 
    [c,  _,  d]] = isMS a d && isMS b c
    where 
        isMS 'M' 'S' = True
        isMS 'S' 'M' = True
        isMS _ _ = False
isXMAS [[_, _, _], [_, _, _], [_, _, _]] = False

part2 :: [String] -> Int
part2 = sum . map (length . filter id) . kernelMap isXMAS Nothing

main :: IO ()
main = aocMain id part1 part2