import pprint
import pandas as pd
import numpy as np

mas_patterns = [
    pd.DataFrame(
        [
            ["M", ".", "M"],
            [".", "A", "."],
            ["S", ".", "S"],
        ]
    ),
    pd.DataFrame(
        [
            ["M", ".", "S"],
            [".", "A", "."],
            ["M", ".", "S"],
        ]
    ),
    pd.DataFrame(
        [
            ["S", ".", "M"],
            [".", "A", "."],
            ["S", ".", "M"],
        ]
    ),
    pd.DataFrame(
        [
            ["S", ".", "S"],
            [".", "A", "."],
            ["M", ".", "M"],
        ]
    ),
]


def get_diagonals(input_df: pd.DataFrame):
    n_cols = len(input_df.columns)
    n_rows = len(input_df)

    diagonals = []

    for i in range(-1 * n_rows + 1, n_cols):
        diagonals.append("".join(np.diag(input_df, i)))

    return diagonals


def get_sub_arrays(input_df: pd.DataFrame):
    n_cols = len(input_df.columns)
    n_rows = len(input_df)

    sub_arrays = []

    for i in range(n_rows - 2):
        for j in range(n_cols - 2):
            sub = input_df.iloc[i : i + 3, j : j + 3].copy()
            sub.iloc[0, 1] = "."
            sub.iloc[1, 0] = "."
            sub.iloc[1, 2] = "."
            sub.iloc[2, 1] = "."
            sub_arrays.append(sub)

    return sub_arrays


with open("./four.txt") as f:
    wordsearch_input = f.readlines()

lines = [list(s.strip()) for s in wordsearch_input if s]
lines = pd.DataFrame(lines)

possible_strings = []

# Rows
possible_strings.extend("".join(line.to_list()) for _, line in lines.iterrows())
possible_strings.extend("".join(line.to_list()) for _, line in lines.items())

# Diagonals tl - br
possible_strings.extend(get_diagonals(lines))

# Diagonals bl - tr
possible_strings.extend(get_diagonals(pd.DataFrame(np.rot90(lines))))

# Find xmas
occurences = 0
for string in possible_strings:
    occurences += string.upper().count("XMAS")
    occurences += string.upper().count("SAMX")

print(occurences)

sub_arrays = get_sub_arrays(lines)


better_occurences = 0
for i, sub_array in enumerate(sub_arrays):
    for j, pattern in enumerate(mas_patterns):
        try:
            equality = (sub_array.values == pattern.values).all()
            if equality:
                better_occurences += 1
                break
        except Exception:
            print(i, j)
            print(sub_array)
            print(pattern)
            raise

print(better_occurences)