--- Day 4: Ceres Search ---

"Looks like the Chief's not here. Next!" One of The Historians pulls out a device and pushes the only button on it. After a brief flash, you recognize the interior of the Ceres monitoring station!

As the search for the Chief continues, a small Elf who lives on the station tugs on your shirt; she'd like to know if you could help her with her word search (your puzzle input). She only has to find one word: XMAS.

This word search allows words to be horizontal, vertical, diagonal, written backwards, or even overlapping other words. It's a little unusual, though, as you don't merely need to find one instance of XMAS - you need to find all of them. Here are a few ways XMAS might appear, where irrelevant characters have been replaced with .:

..X...
.SAMX.
.A..A.
XMAS.S
.X....
The actual word search will be full of letters instead. For example:

MMMSXXMASM
MSAMXMSMSA
AMXSXMAAMM
MSAMASMSMX
XMASAMXAMM
XXAMMXXAMA
SMSMSASXSS
SAXAMASAAA
MAMMMXMMMM
MXMXAXMASX
In this word search, XMAS occurs a total of 18 times; here's the same word search again, but where letters not involved in any XMAS have been replaced with .:

....XXMAS.
.SAMXMS...
...S..A...
..A.A.MS.X
XMASAMX.MM
X.....XA.A
S.S.S.S.SS
.A.A.A.A.A
..M.M.M.MM
.X.X.XMASX
Take a look at the little Elf's word search. How many times does XMAS appear?
from typing import List
import pprint

INPUT_FILE: str = "input.txt"

def readInput() -> List[List[str]]:
    with open(INPUT_FILE, 'r') as f:
        lines = f.readlines()
        return [list(line.strip()) for line in lines]

def search2D(grid, row, col, word):
    # Directions: right, down, left, up, diagonal down-right, diagonal down-left, diagonal up-right, diagonal up-left
    x = [0, 1, 0, -1, 1, 1, -1, -1]
    y = [1, 0, -1, 0, 1, -1, 1, -1]

    lenWord = len(word)
    count = 0

    for dir in range(8):
        k = 0
        currX, currY = row, col

        while k < lenWord:
            if (0 <= currX < len(grid)) and (0 <= currY < len(grid[0])) and (grid[currX][currY] == word[k]):
                currX += x[dir]
                currY += y[dir]
                k += 1
            else:
                break

        if k == lenWord:
            count += 1

    return count

def searchWord(grid, word):
    m = len(grid)
    n = len(grid[0])

    total_count = 0

    for row in range(m):
        for col in range(n):
            total_count += search2D(grid, row, col, word)

    return total_count

def main():
    input = readInput()
    word = "XMAS"
    # pprint.pprint(input)
    result = searchWord(input, word)
    pprint.pprint(result)

if __name__ == "__main__":
    main()