from tools.aoc import AOCDay
from tools.coordinate import Coordinate, DistanceAlgorithm
from tools.grid import Grid
from typing import Any


def get_cheat_count(track: list[Coordinate], check_len: int, need_save: int = 100) -> int:
    total_dist = len(track)
    ans = 0
    for i, c in enumerate(track[:-need_save]):
        for j, n in enumerate(track[i + 2 :]):
            dist = c.getDistanceTo(n, algorithm=DistanceAlgorithm.MANHATTAN, includeDiagonals=False)
            if dist > check_len:
                continue
            saved = total_dist - (i + dist + (total_dist - i - (j + 2)))
            if saved >= need_save:
                ans += 1

    return ans


class Day(AOCDay):
    inputs = [
        [
            (5, "input20_test"),
            (1338, "input20"),
        ],
        [
            (41, "input20_test"),
            (975376, "input20"),
        ],
    ]

    def parse_input(self) -> list[Coordinate]:
        grid = Grid.from_data(self.getInput(), default=False, translate={".": True, "#": False})
        start = list(grid.find("S"))[0]
        end = list(grid.find("E"))[0]

        cur = start
        track = [start]
        while cur != end:
            for n in grid.getNeighboursOf(cur, includeDiagonal=False):
                if n not in track:
                    cur = n
                    track.append(n)

        return track

    def part1(self) -> Any:
        need_save = 20 if self.is_test() else 100
        track = self.parse_input()
        return get_cheat_count(track, need_save=need_save, check_len=2)

    def part2(self) -> Any:
        need_save = 70 if self.is_test() else 100
        track = self.parse_input()
        return get_cheat_count(track, need_save=need_save, check_len=20)


if __name__ == "__main__":
    day = Day(2024, 20)
    day.run(verbose=True)