from typing import Iterable

WORD = "XMAS"
def words_from_grid(grid, startx, starty) -> Iterable[str]:
    num_rows = len(grid)
    num_cols = len(grid[0])

    for (dx, dy) in (
        (1, 0), (-1, 0),
        (0, 1), (0, -1),
        (1, 1), (-1, 1),
        (1, -1), (-1, -1)
    ):
        x = startx
        y = starty
        word = ""
        for i in range(len(WORD)):
            if x >= num_cols or x < 0:
                continue
            if y >= num_rows or y < 0:
                continue

            word += grid[y][x]
            x += dx
            y += dy

        yield word

def cross_words_from_grid(grid, startx, starty) -> tuple[str, str]:
    num_rows = len(grid)
    num_cols = len(grid[0])

    output = []
    for (dx, dy) in (
        (1, 1), (1, -1)
    ):
        x = startx - dx
        y = starty - dy
        word = ""
        for _ in range(3):
            if x >= num_cols or x < 0:
                continue
            if y >= num_rows or y < 0:
                continue

            word += grid[y][x]
            x += dx
            y += dy
        output.append(word)
    # o: tuple[str, str] = 
    return tuple(output)

def task1(grid: list[str]) -> int:
    count = 0
    for y in range(len(grid)):
        for x in range(len(grid[0])):
            for word in words_from_grid(grid, x, y):
                if word == WORD or word == str(reversed(WORD)):
                    count += 1
    return count

def task2(grid: list[str]) -> int:
    count = 0
    for y in range(len(grid)):
        for x in range(len(grid[0])):
            leading, anti = cross_words_from_grid(grid, x, y)
            if (leading == "MAS" or leading == "SAM") and (anti == "MAS" or anti == "SAM"):
                count += 1
    return count

# I can't believe I got both parts first try... this should've been a very error-prone episode!
def main():
    with open("i4.txt") as f:
        text = f.read().splitlines()
    print(task1(text))
    print(task2(text))

if __name__ == "__main__":
    main()