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


class RegionType(int, Enum):
    ROCKY = 0
    WET = 1
    NARROW = 2


class Tools(int, Enum):
    NEITHER = 0
    TORCH = 1
    CLIMBING_GEAR = 2


COMMON_TOOLS = {
    RegionType.ROCKY: {Tools.TORCH, Tools.CLIMBING_GEAR},
    RegionType.WET: {Tools.NEITHER, Tools.CLIMBING_GEAR},
    RegionType.NARROW: {Tools.NEITHER, Tools.TORCH},
}


def get_common_tool(from_region_type: RegionType, to_region_type: RegionType) -> Tools:
    return (COMMON_TOOLS[from_region_type] & COMMON_TOOLS[to_region_type]).pop()


class Cave(Grid):
    def __init__(self, depth: int, target_x: int, target_y: int):
        super().__init__()
        self.depth = depth
        self.target_x = target_x
        self.target_y = target_y
        self.set(Coordinate(0, 0), self.get_erosion_level(Coordinate(0, 0)))

    def get_geologic_index(self, at: Coordinate) -> int:
        if at.x == at.y == 0:
            return 0
        elif at.x == self.target_x and at.y == self.target_y:
            return 0
        elif at.y == 0:
            return at.x * 16807
        elif at.x == 0:
            return at.y * 48271
        else:
            return self.get(at - Coordinate(1, 0)) * self.get(at - Coordinate(0, 1))

    def get_region_type(self, at: Coordinate) -> RegionType:
        return RegionType(self.get(at) % 3)

    def get_erosion_level(self, at: Coordinate) -> int:
        return (self.get_geologic_index(at) + self.depth) % 20183

    def get(self, at: Coordinate) -> int:
        if not self.isSet(at):
            self.set(at, self.get_erosion_level(at))

        return super().get(at)

    def get_risk_level(self) -> int:
        risk = 0
        for x in range(0, self.target_x + 1):
            for y in range(0, self.target_y + 1):
                at = Coordinate(x, y)
                risk += int(self.get_region_type(at))

        return risk


class Day(AOCDay):
    inputs = [
        [
            (114, "input22_test"),
            (8681, "input22"),
        ],
        [
            (45, "input22_test"),
            (1070, "input22"),
        ],
    ]

    def get_params(self) -> (int, int, int):
        """(depth, target_x, target_y)"""
        lines = self.getInput()
        depth = int(lines[0].split()[1])
        target_x, target_y = map(int, lines[1].split()[1].split(","))
        return depth, target_x, target_y

    def part1(self) -> Any:
        cave = Cave(*self.get_params())
        return cave.get_risk_level()

    def part2(self) -> Any:
        params = self.get_params()
        target = Coordinate(*params[1:])
        cave = Cave(*params)
        queue = [(0, 0, Tools.TORCH, Coordinate(0, 0))]
        heapify(queue)
        visited = set()

        target_dist, target_tool = None, None
        while len(queue) > 0:
            _, current_dist, current_tool, at = heappop(queue)
            if (at, current_tool) in visited:
                continue

            visited.add((at, current_tool))

            if at == target:
                target_dist = current_dist
                target_tool = current_tool
                break

            for next in at.getNeighbours(includeDiagonal=False, minX=0, minY=0):
                region_type_at = cave.get_region_type(at)
                region_type_next = cave.get_region_type(next)
                if region_type_at == region_type_next:
                    common_tool = current_tool
                else:
                    common_tool = get_common_tool(region_type_at, region_type_next)

                if common_tool == current_tool:
                    heappush(
                        queue, (next.getDistanceTo(target) + current_dist + 1, current_dist + 1, common_tool, next)
                    )
                else:
                    heappush(
                        queue, (next.getDistanceTo(target) + current_dist + 7, current_dist + 8, common_tool, next)
                    )

        if target_tool != Tools.TORCH:
            target_dist += 7
            while queue:
                _, current_dist, current_tool, at = heappop(queue)
                if at != target or current_dist >= target_dist:
                    continue

                if current_tool == Tools.TORCH:
                    target_dist = current_dist
                elif current_dist + 7 < target_dist:
                    target_dist = current_dist + 7

        return target_dist


if __name__ == "__main__":
    day = Day(2018, 22)
    day.run(verbose=True)