from heapq import heappush, heappop
from tools.aoc import AOCDay
from tools.coordinate import Coordinate
from tools.grid import Grid
from typing import Any


def expand_valley(valley: Grid, to_z: int) -> None:
    while valley.maxZ <= to_z:
        for blizzard_pos in valley.getActiveCells(z=valley.maxZ):
            for blizzard in valley.get(blizzard_pos):
                next_pos = blizzard_pos + blizzard
                if next_pos.x < 0:
                    next_pos = Coordinate(valley.maxX, next_pos.y, next_pos.z)
                elif next_pos.x > valley.maxX:
                    next_pos = Coordinate(0, next_pos.y, next_pos.z)
                elif next_pos.y < 0:
                    next_pos = Coordinate(next_pos.x, valley.maxY, next_pos.z)
                elif next_pos.y > valley.maxY:
                    next_pos = Coordinate(next_pos.x, 0, next_pos.z)

                if not valley.get(next_pos):
                    valley.set(next_pos, [blizzard])
                else:
                    valley.get(next_pos).append(blizzard)


def get_min_steps(valley: Grid, pos: Coordinate, target: Coordinate) -> int:
    check_list = [
        Coordinate(1, 0, 1),
        Coordinate(-1, 0, 1),
        Coordinate(0, 0, 1),
        Coordinate(0, 1, 1),
        Coordinate(0, -1, 1)
    ]
    q = []
    heappush(q, (0, pos))
    v = set()
    while q:
        _, cur_pos = heappop(q)
        if cur_pos in v:
            continue
        v.add(cur_pos)

        if cur_pos.z >= valley.maxZ:
            expand_valley(valley, cur_pos.z + 1)

        for c in check_list:
            next_pos = cur_pos + c
            if next_pos.x == target.x and next_pos.y == target.y:
                return next_pos.z
            if not valley.isWithinBoundaries(next_pos) and (next_pos.x != pos.x or next_pos.y != pos.y):
                continue
            if valley.get(next_pos):
                continue
            dist = abs(target.x - next_pos.x) + abs(target.y - next_pos.y)
            heappush(q, (dist + next_pos.z, next_pos))

    return -1


class Day(AOCDay):
    inputs = [
        [
            (18, "input24_test"),
            (297, "input24_dennis"),
            (269, "input24"),
        ],
        [
            (54, "input24_test"),
            (856, "input24_dennis"),
            (825, "input24"),
        ]
    ]

    def get_valley(self) -> Grid:
        valley = Grid.from_str(
            "/".join(self.getInput()),
            translate={
                '#': False,
                '.': False,
                '>': [Coordinate(1, 0, 1)],
                '<': [Coordinate(-1, 0, 1)],
                'v': [Coordinate(0, 1, 1)],
                '^': [Coordinate(0, -1, 1)],
            },
            default=False,
            mode3d=True
        )
        valley.shift_zero()

        return valley

    def part1(self) -> Any:
        valley = self.get_valley()
        return get_min_steps(valley, Coordinate(0, -1, 0), Coordinate(valley.maxX, valley.maxY + 1))

    def part2(self) -> Any:
        valley = self.get_valley()
        start = Coordinate(0, -1, 0)
        target = Coordinate(valley.maxX, valley.maxY + 1)
        steps = get_min_steps(valley, start, target)
        steps = get_min_steps(valley, Coordinate(target.x, target.y, steps), start)
        return get_min_steps(valley, Coordinate(start.x, start.y, steps), target)


if __name__ == '__main__':
    day = Day(2022, 24)
    day.run(verbose=True)