#!/usr/bin/env python3

with open("./2024/inputs/day04.txt") as f:
    puzzle = [l.strip() for l in f.readlines()]

maxY = len(puzzle) - 1
maxX = len(puzzle[0]) - 1

def possible_words(p, start):
    y = start[0]
    x = start[1]
    ret = []

    # get the two possible horizontal matches
    # skip them if x - 3 < 0 (word would be truncated on the left)
    # or x + 3 > maxX (word would be truncated on the right)
    if not ((x + 3) > maxX):
        ret.append(str(f"{p[y][x]}{p[y][x + 1]}{p[y][x + 2]}{p[y][x + 3]}"))
    
    if not ((x - 3) < 0):
        ret.append(str(f"{p[y][x]}{p[y][x - 1]}{p[y][x - 2]}{p[y][x - 3]}"))
    
    # get the two possible vertical matches
    # skip them if y - 3 < 0 (word would be truncated on the top)
    # or y + 3 > maxY (word would be truncated on the bottom)
    if not ((y + 3) > maxY):
        ret.append(str(f"{p[y][x]}{p[y + 1][x]}{p[y + 2][x]}{p[y + 3][x]}"))
    
    if not ((y - 3) < 0):
        ret.append(str(f"{p[y][x]}{p[y - 1][x]}{p[y - 2][x]}{p[y - 3][x]}"))
    
    # get the four diagonals, again with some (basic) bounds checking
    if not ((y + 3) > maxY) and not ((x + 3) > maxX):
        ret.append(str(f"{p[y][x]}{p[y + 1][x + 1]}{p[y + 2][x + 2]}{p[y + 3][x + 3]}"))
    
    if not ((y - 3) < 0) and not ((x - 3) < 0):
        ret.append(str(f"{p[y][x]}{p[y - 1][x - 1]}{p[y - 2][x - 2]}{p[y - 3][x - 3]}"))
    
    if not ((y + 3) > maxY) and not ((x - 3) < 0):
        ret.append(str(f"{p[y][x]}{p[y + 1][x - 1]}{p[y + 2][x - 2]}{p[y + 3][x - 3]}"))
    
    if not ((y - 3) < 0) and not ((x + 3) > maxX):
        ret.append(str(f"{p[y][x]}{p[y - 1][x + 1]}{p[y - 2][x + 2]}{p[y - 3][x + 3]}"))
    
    return ret

xmas_count = 0
for y in range(maxY):
    for x in range(maxX):
        if puzzle[y][x] == "X":
            words = possible_words(puzzle, (y, x))
            for word in words:
                xmas_count += word.count("XMAS")
    
        if puzzle[y][x] == "S":
            words = possible_words(puzzle, (y, x))
            for word in words:
                xmas_count += word.count("SAMX")

print(xmas_count)