--- 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 re import findall

with open("input.txt") as input_file:
    input_text = input_file.read().splitlines()

num_rows = len(input_text)
num_cols = len(input_text[0])
total = 0

# Rows
for row in input_text:
    total += len(findall(r"XMAS", row))
    total += len(findall(r"XMAS", row[::-1]))

# Columns
for col_idx in range(num_cols):
    col = "".join([input_text[row_idx][col_idx] for row_idx in range(num_rows)])
    total += len(findall(r"XMAS", col))
    total += len(findall(r"XMAS", col[::-1]))
# NE/SW Diagonals
for idx_sum in range(num_rows + num_cols - 1):
    diagonal = "".join(
        [
            input_text[row_idx][idx_sum - row_idx]
            for row_idx in range(
                max(0, idx_sum - num_cols + 1), min(num_rows, idx_sum + 1)
            )
        ]
    )
    total += len(findall(r"XMAS", diagonal))
    total += len(findall(r"XMAS", diagonal[::-1]))
# NW/SE Diagonals
for idx_diff in range(-num_cols + 1, num_rows):
    diagonal = "".join(
        [
            input_text[row_idx][row_idx - idx_diff]
            for row_idx in range(max(0, idx_diff), min(num_rows, num_cols + idx_diff))
        ]
    )
    total += len(findall(r"XMAS", diagonal))
    total += len(findall(r"XMAS", diagonal[::-1]))
print(total)