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

kernelMapRow :: ([[a]] -> b) -> [a] -> [a] -> [a] -> [b]
kernelMapRow kernel
    (a:r0@(b:c:_))
    (d:r1@(e:f:_))
    (g:r2@(h:i:_)) =
         kernel [[a,b,c],[d,e,f],[g,h,i]] : kernelMapRow kernel r0 r1 r2
kernelMapRow _ [_, _] [_, _] [_, _] = []

kernelMap :: ([[a]] -> b) -> [[a]] -> [[b]]
kernelMap kernel (a:rest@(b:c:_)) = kernelMapRow kernel a b c : kernelMap kernel rest
kernelMap _ [_, _] = []

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

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