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


def getGrid(lines: List[str], multiply: bool = False) -> Grid:
    size = len(lines)
    g = Grid()
    for y, l in enumerate(lines):
        for x, v in enumerate(map(int, l)):
            g.set(Coordinate(x, y), v)
            if not multiply:
                continue

            for x2 in range(5):
                for y2 in range(5):
                    if x2 == y2 == 0:
                        continue
                    nv = 1 + (v + x2 + y2 - 1) % 9

                    g.set(Coordinate(size * x2 + x, size * y2 + y), nv)

    return g


class Day(AOCDay):
    test_solutions_p1 = [40, 386]
    test_solutions_p2 = [315, 2806]

    def part1(self) -> Any:
        grid = getGrid(self.getInput())
        path = grid.getPath(Coordinate(0, 0), Coordinate(grid.maxX, grid.maxY), includeDiagonal=False, weighted=True)
        return sum(grid.get(c) for c in path[:-1])

    def part2(self) -> Any:
        grid = getGrid(self.getInput(), True)
        path = grid.getPath(Coordinate(0, 0), Coordinate(grid.maxX, grid.maxY), includeDiagonal=False, weighted=True)
        return sum(grid.get(c) for c in path[:-1])