import           Data.List (isPrefixOf, transpose)
import           Data.Universe.Helpers (diagonals)

countTimes :: String -> String -> Int
countTimes _ [] = 0
countTimes s fullLine@(_:nextLine)
  | s `isPrefixOf` fullLine = 1 + countTimes s (drop (length s) fullLine)
  | otherwise = countTimes s nextLine

countHoriz :: String -> Int
countHoriz l = countTimes "XMAS" l + countTimes "SAMX" l

countLines :: [String] -> Int
countLines = sum . map countHoriz

countAll :: [String] -> Int
countAll ls = countLines ls
  + countLines (transpose ls)
  + countLines (diagonals ls)
  + countLines (diagonals (reverse ls))

main :: IO ()
main = print . countAll . lines =<< readFile "data.txt"