#!/usr/bin/env python
# coding: utf-8

# # Day 22 - More pathfinding, adding weights
# 
# - [Day 22](https://adventofcode.com/2018/day/22)
# 
# At first this appears to be a matrix problem, but the modulus over 3 prime numbers rules that out quickly. Part two reveals that this is actually a pathfinding problem, with part 1 just setting you up with a map generator.
# 
# The map is essentially infinite in the positive `x` and `y` directions, so you need to generate more map as you try to find your path. You could add a row or column as you reach a new highest `x` or `y` position, but those function calls and list appends will add up to a rather high amount of time.
# 
# So in my solution I double the row or column count each time a coordinate strays outside of the generated erosion level map, amortising the cost of expanding; the function only needs to be called a few times for the full rescue operation.
# 

# In[1]:


import math
from dataclasses import dataclass
from enum import Enum
from heapq import heappop, heappush
from itertools import accumulate, chain, count, islice
from typing import Iterator, List, Tuple

M: int = 20183
X0: int = 16807
Y0: int = 48271

Pos = Tuple[int, int]


class Tool(Enum):
    neither = 0
    torch = 1
    climbing_gear = 2


class RegionType(Enum):
    rocky = 0, ".", Tool.climbing_gear, Tool.torch
    wet = 1, "=", Tool.climbing_gear, Tool.neither
    narrow = 2, "|", Tool.torch, Tool.neither

    def __new__(cls, value, char, *tools):
        instance = object.__new__(cls)
        instance._value_ = value
        instance.char = char
        instance.tools = tools
        return instance

    def __str__(self) -> str:
        return self.char


@dataclass(frozen=True)
class Node:
    """Node on the A* search graph"""

    x: int = 0
    y: int = 0
    tool: Tool = Tool.torch
    time: int = 0

    @property
    def pos(self) -> Pos:
        return self.x, self.y

    def cost(self, target: Pos) -> int:
        """Calculate the cost for this node, f(n) = g(n) + h(n)

        The cost of this node is the time taken (g) plus
        estimated cost to get to nearest goal (h).

        Here we use the manhattan distance to the target as
        the estimated cost.

        """
        return self.time + abs(target[0] - self.x) + abs(target[1] - self.y)

    def transitions(self, maze: "ModeMaze") -> Iterator["Node"]:
        cls = type(self)
        other = next(t for t in maze[self.pos].tools if t is not self.tool)
        # switching tools
        yield cls(self.x, self.y, other, self.time + 7)
        positions = (
            (self.x + dx, self.y + dy) for dx, dy in ((-1, 0), (0, -1), (0, 1), (1, 0))
        )
        yield from (
            cls(x, y, self.tool, self.time + 1)
            for x, y in positions
            if x >= 0 and y >= 0 and self.tool in maze[x, y].tools
        )


class ModeMaze:
    def __init__(self, depth: int, target: Pos) -> None:
        self.depth = depth
        self.target = target
        # y rows, x columns
        self._erosion_levels: List[List[int]] = [[0]]
        self._width, self._height = 1, 1

    def _gen_levels(self, pos: Pos) -> None:
        # make sure there is enough erosion level info up to the given position
        # We'll grow this dynamically to the next power of 2 larger than needed each time
        # to amortize the cost
        cw, ch = self._width, self._height
        width = max(2 ** math.ceil(math.log2(pos[0] + 1)), cw)
        height = max(2 ** math.ceil(math.log2(pos[1] + 1)), ch)

        depth = self.depth
        target = self.target
        maze = self._erosion_levels

        # accumulator function
        if cw <= target[1] < width or ch <= target[0] < height:
            # takes a nonlocal cx as a counter
            cx = None

            def geoindex(px, py, m=M):
                return depth if (next(cx), y) == target else (px * py + depth) % m

        else:

            def geoindex(px, py, m=M):
                return (px * py + depth) % m

        if width > cw:
            # add more columns
            rows = iter(maze)
            next(rows).extend((x * X0 + depth) % M for x in range(cw, width))
            # iterate over rows pairwise, so we get the preceding and current row together
            for p, row in zip(maze, rows):
                # accumulate values based on the preceding column value and the row above
                cx = count(cw)  # only used if target position is in range
                new_x = accumulate(
                    chain(row[-1:], islice(p, len(row), width)), geoindex
                )
                next(new_x)  # this is the last value in the row, row[-1]
                row += new_x

        if height > ch:
            # add more rows; pairwise iteration with the preceding row (last in the maze so far)
            for y, p in zip(range(ch, height), islice(maze, ch - 1, None)):
                row = [(y * Y0 + self.depth) % M]
                cx = count(1)
                row = list(accumulate(chain(row, islice(p, 1, width)), geoindex))
                maze.append(row)

        self._width, self._height = width, height

    @property
    def risk_level(self) -> int:
        assert self[self.target] is RegionType.rocky
        tx, ty = self.target
        return sum(
            sum(i % 3 for i in row[: tx + 1]) for row in self._erosion_levels[: ty + 1]
        )

    def __getitem__(self, pos: Pos) -> RegionType:
        if pos[0] >= self._width or pos[1] >= self._height:
            self._gen_levels(pos)
        return RegionType(self._erosion_levels[pos[1]][pos[0]] % 3)

    def __str__(self) -> str:
        lines = [[str(RegionType(i % 3)) for i in row] for row in self._erosion_levels]
        try:
            lines[0][0] = "M"
            lines[self.target[1]][self.target[0]] = "T"
        except IndexError:
            pass
        return "\n".join(["".join(line) for line in lines])

    def rescue(self) -> int:
        start = Node()
        open = {start}
        unique = count()  # tie breaker when costs are equal
        pqueue = [(start.cost(self.target), next(unique), start)]
        closed = set()
        times = {
            (start.pos, start.tool): start.time
        }  # (pos, tool) -> time. Ignore nodes that take longer
        while open:
            node = heappop(pqueue)[-1]

            if node.pos == self.target and node.tool is Tool.torch:
                return node.time

            open.remove(node)
            closed.add(node)
            for new in node.transitions(self):
                if new in closed or new in open:
                    continue
                if times.get((new.pos, new.tool), float("inf")) < new.time:
                    continue
                times[new.pos, new.tool] = new.time
                open.add(new)
                heappush(pqueue, (new.cost(self.target), next(unique), new))


# In[2]:


test = ModeMaze(510, (10, 10))
assert test.risk_level == 114
assert (
    str(test)
    == """\
M=.|=.|.|=.|=|=.
.|=|=|||..|.=...
.==|....||=..|==
=.|....|.==.|==.
=|..==...=.|==..
=||.=.=||=|=..|=
|.=.===|||..=..|
|..==||=.|==|===
.=..===..=|.|||.
.======|||=|=.|=
.===|=|===T===||
=|||...|==..|=.|
=.=|=.=..=.||==|
||=|=...|==.=|==
|=.=||===.|||===
||.|==.|.|.||=||"""
)
assert test.rescue() == 45


# In[3]:


import re

import aocd

data = aocd.get_data(day=22, year=2018)
depth, tx, ty = map(int, re.findall(r"\d+", data))
maze = ModeMaze(depth, (tx, ty))


# In[4]:


print("Part 1:", maze.risk_level)


# In[5]:


print("Part 2:", maze.rescue())