import numpy as np
from pathlib import Path

filepath = Path("./data/4_example")
filepath = Path("./data/4_input")

directions = np.array([
    [[0,0], [1,0], [2,0], [3,0]],
    [[0,0], [-1,0], [-2,0], [-3,0]],
    [[0,0], [0,1], [0,2], [0,3]],
    [[0,0], [0,-1], [0,-2], [0,-3]],
    [[0,0], [1,1], [2,2], [3,3]],
    [[0,0], [-1,1], [-2,2], [-3,3]],
    [[0,0], [1,-1], [2,-2], [3,-3]],
    [[0,0], [-1,-1], [-2,-2], [-3,-3]],
])

with open(filepath) as filein:
    data = [line.rstrip() for line in filein.readlines()]

def find_xmas(arr, i,j):
    count = 0
    for d in directions:
        x,m,a,s = "0000"
        try:
            assert np.all([i,j] + d >= 0)
            x = arr[i][j]
            m = arr[i+d[1,0]][j+d[1,1]] 
            a = arr[i+d[2,0]][j+d[2,1]] 
            s = arr[i+d[3,0]][j+d[3,1]] 
            if x+m+a+s == "XMAS":
                count += 1
        except IndexError:
            pass
        except AssertionError:
            pass
    return count

global_count = 0
for i in range(len(data)):
    for j in range(len(data[0])):
        if data[i][j] == 'X':
            count = find_xmas(data, i, j)
            global_count += count
    #         print(count if count else '.', end="")
    #     else:
    #         print(".", end="")
    # print('')

print(global_count)


data = np.array([[c for c in line] for line in data])

directions = np.array([
    [[-1,-1],[0,0],[1,1]],
    [[1,1],[0,0],[-1,-1]],
    [[-1,1],[0,0],[1,-1]],
    [[1,-1],[0,0],[-1,1]],
])

def find_x_mas(arr, i,j):
    count = 0
    for dir in directions:
        try:
            d = dir + [i,j]
            assert np.all(d >= 0)
            m = arr[d[0,0],d[0,1]]
            a = arr[d[1,0],d[1,1]]
            s = arr[d[2,0],d[2,1]]
            if m+a+s == "MAS":
                count += 1
        except IndexError:
            pass
        except AssertionError:
            pass
    if count == 2:
        return True
    else:
        return 0

global_count = 0
n,m = data.shape
for i in range(n):
    for j in range(m):
        if data[i][j] == 'A':
            count = find_x_mas(data, i, j)
            global_count += count
    #         print(count if count else '.', end="")
    #     else:
    #         print(".", end="")
    # print('')

print(global_count)