import aocd

from enum import Enum, nonmember
import collections

class Direction(Enum):
    NORTH = (0, -1)
    NORTHEAST = (1, -1)
    EAST = (1, 0)
    SOUTHEAST = (1, 1)
    SOUTH = (0, 1)
    SOUTHWEST = (-1, 1)
    WEST = (-1, 0)
    NORTHWEST = (-1, -1)

    @nonmember
    def opposite(d):
        return Direction((d.value[0]*-1, d.value[1]*-1))

class Letter:
    target = 'MAS'

    def __init__(self, char, x, y, grid):
        self.char = char

        self.x = x
        self.y = y
        self.grid = grid
    
    def neighbour_coords(self, direction):
        return (self.x + direction.value[0], self.y + direction.value[1])

    def neighbour(self, direction):
        return self.grid[self.neighbour_coords(direction)]

    def is_center(self):
        """Return whether this letter is in the center of the target pattern"""
        if self.char != self.target[1]:
            return False
        
        corners = (
            Direction.NORTHWEST,
            Direction.NORTHEAST,
            Direction.SOUTHWEST,
            Direction.SOUTHEAST,
        )

        return sum(self.neighbour(d)._is_nth_position(0, Direction.opposite(d)) for d in corners) == 2

    def _is_nth_position(self, n, direction):
        """
        Return number of times this letter is the nth letter of the target word 
        in the given direction
        """
        # Must be the correct letter
        if self.char != self.target[n]:
            return False
        
        # Return True if it's the last letter
        if n == len(self.target) - 1:
            return True

        # Return False if there are no more letter in this direction
        if (next_letter := self.neighbour(direction)) is None:
            return False
        
        # Ask the next letter
        return next_letter._is_nth_position(n+1, direction)

def parse(data):
    """
    Parse string into grid, return all possible first letters of XMAS
    """
    grid = collections.defaultdict(lambda: Letter(None, None, None, None)) # Null letters
    a_list = []

    for y, r in enumerate(data.split('\n')):
        for x, c in enumerate(r):
            grid[(x, y)] = Letter(c, x, y, grid)
            if c == 'A':
                a_list.append(grid[(x, y)])
    
    return a_list

def solve(data):
    a_list = parse(data)
    return sum(a.is_center() for a in a_list)

if __name__ == '__main__':
    solution = solve(aocd.get_data(year=2024, day=4))

    print(solution)
    aocd.submit(solution, year=2024, day=4, part='b')