from collections import deque, defaultdict
from tools.aoc import AOCDay
from tools.coordinate import Coordinate
from tools.grid import Grid
from typing import Any


def walk(grid: Grid, start: Coordinate, steps: int) -> int:
    q = deque()
    q.append((start, 0))
    s_mod = steps % 2
    v = set()
    v2 = set()
    while q:
        coord, dist = q.popleft()
        #print("check ", coord, dist)
        if coord in v or coord in v2:
            #print(" -> in v")
            continue
        if dist % 2 == s_mod:
            #print(f" -> add to v because {dist} even or == {steps}")
            v.add(coord)
        else:
            v2.add(coord)
        if dist >= steps:
            #print(f" -> >= {steps}")
            continue
        for n in coord.getNeighbours(includeDiagonal=False):
            if grid.get(n % (grid.maxX + 1)):
                q.append((n, dist + 1))

    #print([x for x in v])
    return len(v)


class Day(AOCDay):
    inputs = [
        [
            (16, "input21_test"),
            (3562, "input21"),
        ],
        [
            #(16733044, "input21_test"),
            (None, "input21"),
        ]
    ]

    def parse_input(self) -> (Grid, Coordinate):
        grid = Grid().from_data(self.getInput(), translate={'.': True, '#': False})
        start = None
        for x in grid.getActiveCells():
            if grid.get(x) == 'S':
                start = x
                grid.set(x, True)

        return grid, start

    def part1(self) -> Any:
        needed_steps = 6 if 'test' in self._current_test_file else 64
        grid, start = self.parse_input()
        return walk(grid, start, needed_steps)

    def part2(self) -> Any:
        grid, start = self.parse_input()
        needed_steps = 5000 if 'test' in self._current_test_file else 26501365
        k = walk(grid, start, len(self.getInput()))
        g = walk(grid, start, 2 * len(self.getInput()))
        v = walk(grid, start, 3 * len(self.getInput()))
        print(k, g, v, g-k, v-g)
        return ""
        j = 0
        mem = defaultdict(list)
        for i in range(1, 250, 2):
            k = walk(grid, start, i)
            diff = k - j
            j = k
            mem[diff].append(i)
            print(i, "=>", k, "diff", diff, "mod", k % 81)
        for k, v in mem.items():
            if len(v) > 1:
                print(k, "=>", v)
        return ""


if __name__ == '__main__':
    day = Day(2023, 21)
    day.run(verbose=True)