{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Try not to step in the puddles too much\n",
    "\n",
    "- https://adventofcode.com/2023/day/17\n",
    "\n",
    "As has become an Adent of Code tradition, we've been giving a path solving game. The twist is in figuring out what constitutes a 'state' in the path.\n",
    "\n",
    "At first glance you'd give each position in the grid a state; you'd track the direction, the position, and the number of steps in a single direction. However, the number of steps used to reach a position is actually not all that relevant for any position that then changes direction on the next 'move'; for an equal total heat loss it doesn't matter if it took you 1, 2 or 3 same-direction steps to reach a specific position. Or even if you came from one, or the opposite direction.\n",
    "\n",
    "Let me illustrate. Here is an example section of grid, and in the middle, at position D-4, is a crucible that moved down into that square, and it'll turn either left or right at the next move:\n",
    "\n",
    "|     | A   | B   | C   | D   | E   | F   | G   |\n",
    "| --- | --- | --- | --- | --- | --- | --- | --- |\n",
    "| 1   |     |     |     |     |     |     |     |\n",
    "| 2   |     |     |     |     |     |     |     |\n",
    "| 3   |     |     |     |     |     |     |     |\n",
    "| 4   |     |     |     | V   |     |     |     |\n",
    "| 5   |     |     |     |     |     |     |     |\n",
    "| 6   |     |     |     |     |     |     |     |\n",
    "| 7   |     |     |     |     |     |     |     |\n",
    "\n",
    "Because it can only move up to 3 steps in a single direction, it could reach any of the cells along the row, so A-4, B-4, C-4, E-4, F-4 or G-4. It can't go further than that, it will have to turn on one of those 6 positions, and it then can only go up or down. Let's mark those positions with `|` bars:\n",
    "\n",
    "|     | A   | B   | C   | D   | E   | F   | G   |\n",
    "| --- | --- | --- | --- | --- | --- | --- | --- |\n",
    "| 1   |     |     |     |     |     |     |     |\n",
    "| 2   |     |     |     |     |     |     |     |\n",
    "| 3   |     |     |     |     |     |     |     |\n",
    "| 4   | \\|  | \\|  | \\|  | V   | \\|  | \\|  | \\|  |\n",
    "| 5   |     |     |     |     |     |     |     |\n",
    "| 6   |     |     |     |     |     |     |     |\n",
    "| 7   |     |     |     |     |     |     |     |\n",
    "\n",
    "Now, we can already see that if the crucible had reached this position from below (moving up), the same 6 positions could have been reached.\n",
    "\n",
    "If the crucible had reached this position from the left or right, then the crucible can only move up or down; it's the same picture as before, but rotated by 90 degrees. Because it doesn't matter if the crucible moved left or right, I've marked the position in the following table with `|` again, and all the positions it can reach with `-`:\n",
    "\n",
    "|     | A   | B   | C   | D   | E   | F   | G   |\n",
    "| --- | --- | --- | --- | --- | --- | --- | --- |\n",
    "| 1   |     |     |     | -   |     |     |     |\n",
    "| 2   |     |     |     | -   |     |     |     |\n",
    "| 3   |     |     |     | -   |     |     |     |\n",
    "| 4   |     |     |     | \\|  |     |     |     |\n",
    "| 5   |     |     |     | -   |     |     |     |\n",
    "| 6   |     |     |     | -   |     |     |     |\n",
    "| 7   |     |     |     | -   |     |     |     |\n",
    "\n",
    "So, for this path finding problem, the individual step states track position and orientation, and we can then generate the next set of steps by making those 6 steps in the two directions and flipping the orientation.\n",
    "\n",
    "Next, I'm using [A\\* search](https://en.wikipedia.org/wiki/A*_search_algorithm) for this problem and so need a _heuristic_ to prioritise nodes more likely to reach the goal with minimal heat loss. The heuristic must, at most, be equal to the actual final cost of the path. I'm using the manhattan distance between the position and the end goal as the lower bound heat loss, because every cell in the map will incur at least 1 unit of loss. Then, simply add the heat loss of the path so far and we have a great heuristic!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import annotations\n",
    "\n",
    "import typing as t\n",
    "from dataclasses import dataclass, field\n",
    "from enum import IntEnum\n",
    "from heapq import heapify, heappop, heappush\n",
    "from itertools import count\n",
    "\n",
    "\n",
    "class PriorityQueue[T]:\n",
    "    def __init__(self, *initial: tuple[int, T]) -> None:\n",
    "        self._queue: list[tuple[int, int, T]] = []\n",
    "        self._count = count()\n",
    "        for pri, item in initial:\n",
    "            self.put(pri, item)\n",
    "        heapify(self._queue)\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self._queue)\n",
    "\n",
    "    def put(self, pri: int, item: T) -> None:\n",
    "        heappush(self._queue, (pri, next(self._count), item))\n",
    "\n",
    "    def get(self) -> T:\n",
    "        if not self:\n",
    "            raise ValueError(\"Queue is empty\")\n",
    "        return heappop(self._queue)[-1]\n",
    "\n",
    "\n",
    "class Pos(t.NamedTuple):\n",
    "    x: int = 0\n",
    "    y: int = 0\n",
    "\n",
    "    def __add__(  # pyright: ignore[reportIncompatibleMethodOverride]\n",
    "        self, other: t.Self\n",
    "    ) -> t.Self:\n",
    "        if isinstance(other, Pos):\n",
    "            return self._replace(x=self.x + other.x, y=self.y + other.y)\n",
    "        return NotImplemented\n",
    "\n",
    "    def __sub__(self, other: t.Self) -> int:\n",
    "        if isinstance(other, Pos):\n",
    "            return abs(self.x - other.x) + abs(self.y - other.y)\n",
    "        return NotImplemented\n",
    "\n",
    "\n",
    "class Orientation(IntEnum):\n",
    "    # value, -delta, +delta, orthogonal\n",
    "    hor = 1, Pos(0, -1), Pos(0, 1), \"ver\"\n",
    "    ver = 2, Pos(-1, 0), Pos(1, 0), \"hor\"\n",
    "\n",
    "    if t.TYPE_CHECKING:\n",
    "        negd: Pos\n",
    "        posd: Pos\n",
    "        _opposite: str\n",
    "\n",
    "    else:\n",
    "\n",
    "        def __new__(\n",
    "            cls, value: int, negd: Pos, posd: Pos, orthogonal: str\n",
    "        ) -> Direction:\n",
    "            inst = int.__new__(cls, value)\n",
    "            inst._value_ = value\n",
    "            inst.negd, inst.posd, inst._opposite = negd, posd, orthogonal\n",
    "            return inst\n",
    "\n",
    "    @property\n",
    "    def orthogonal(self) -> Orientation:\n",
    "        return Orientation[self._opposite]\n",
    "\n",
    "\n",
    "@dataclass(frozen=True, slots=True)\n",
    "class CruciblePosition:\n",
    "    pos: Pos\n",
    "    orientation: Orientation\n",
    "    heat_loss: int = field(compare=False, default=0)\n",
    "\n",
    "    def heuristic(self, goal: Pos) -> int:\n",
    "        return self.heat_loss + (self.pos - goal)\n",
    "\n",
    "    def moves(self, map: Map) -> t.Iterable[t.Self]:\n",
    "        orient = self.orientation\n",
    "        orthogonal = orient.orthogonal\n",
    "        new = type(self)\n",
    "\n",
    "        pos, hloss = self.pos, self.heat_loss\n",
    "        for _ in range(3):\n",
    "            pos += orient.negd\n",
    "            if (loss := map.get(pos)) is None:\n",
    "                break\n",
    "            hloss += loss\n",
    "            yield new(pos, orthogonal, hloss)\n",
    "\n",
    "        pos, hloss = self.pos, self.heat_loss\n",
    "        for _ in range(3):\n",
    "            pos += orient.posd\n",
    "            if (loss := map.get(pos)) is None:\n",
    "                break\n",
    "            hloss += loss\n",
    "            yield new(pos, orthogonal, hloss)\n",
    "\n",
    "\n",
    "class Map(t.Mapping[Pos, int]):\n",
    "    _map: dict[Pos, int]\n",
    "    _goal: Pos\n",
    "\n",
    "    def __init__(self, map_text: str) -> None:\n",
    "        lines = map_text.splitlines()\n",
    "        self._map = {\n",
    "            Pos(x, y): int(hl)\n",
    "            for y, line in enumerate(lines)\n",
    "            for x, hl in enumerate(line)\n",
    "        }\n",
    "        self._goal = Pos(len(lines[0]) - 1, len(lines) - 1)\n",
    "\n",
    "    def __getitem__(self, pos: Pos) -> int:\n",
    "        return self._map[pos]\n",
    "\n",
    "    def __iter__(self) -> t.Iterator[Pos]:\n",
    "        return iter(self._map)\n",
    "\n",
    "    def __len__(self) -> int:\n",
    "        return len(self._map)\n",
    "\n",
    "    def __contains__(self, pos: object) -> bool:\n",
    "        return pos in self._map\n",
    "\n",
    "    def least_heat_loss(\n",
    "        self, crucible_type: type[CruciblePosition] = CruciblePosition\n",
    "    ) -> int:\n",
    "        goal = self._goal\n",
    "        start = Pos()\n",
    "        open: dict[CruciblePosition, int] = {\n",
    "            crucible_type(start, Orientation.hor): 0,\n",
    "            crucible_type(start, Orientation.ver): 0,\n",
    "        }\n",
    "        queue: PriorityQueue[CruciblePosition] = PriorityQueue(\n",
    "            *((cp.heuristic(goal), cp) for cp in open)\n",
    "        )\n",
    "        closed: set[CruciblePosition] = set()\n",
    "\n",
    "        while open:\n",
    "            current = queue.get()\n",
    "\n",
    "            if open.get(current) != current.heat_loss:\n",
    "                # a better path was found already\n",
    "                continue\n",
    "\n",
    "            if current.pos == goal:\n",
    "                return current.heat_loss\n",
    "\n",
    "            del open[current]\n",
    "            closed.add(current)\n",
    "            for neighbor in current.moves(self):\n",
    "                if neighbor in closed:\n",
    "                    continue\n",
    "                if open.get(neighbor, float(\"inf\")) <= neighbor.heat_loss:\n",
    "                    continue\n",
    "                open[neighbor] = neighbor.heat_loss\n",
    "                queue.put(neighbor.heuristic(goal), neighbor)\n",
    "\n",
    "        raise AssertionError(\"should never reach here\")\n",
    "\n",
    "\n",
    "test_input = \"\"\"\\\n",
    "2413432311323\n",
    "3215453535623\n",
    "3255245654254\n",
    "3446585845452\n",
    "4546657867536\n",
    "1438598798454\n",
    "4457876987766\n",
    "3637877979653\n",
    "4654967986887\n",
    "4564679986453\n",
    "1224686865563\n",
    "2546548887735\n",
    "4322674655533\n",
    "\"\"\"\n",
    "test_map = Map(test_input)\n",
    "assert test_map.least_heat_loss() == 102"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Part 1: 1256\n"
     ]
    }
   ],
   "source": [
    "import aocd\n",
    "\n",
    "map = Map(aocd.get_data(day=17, year=2023))\n",
    "print(\"Part 1:\", map.least_heat_loss())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Taking bigger steps, further\n",
    "\n",
    "With part two now revealed, I am _very_ pleased with my choice of states for the path search! Because the ultra crucible path is now simply the set of positions between 4 and 10 steps away. We never have to worry about wether or not the crucible has made enough steps to stop at teh end goal here, because a state will never produce a next move that's fewer than 4 steps away.\n",
    "\n",
    "With A\\* search pruning the search tree and not having to track states for every single step and direction, part 2 is solved very quickly, in about a second. I also tried a straight up Dijkstra (breath-first) search, without the heuristic, and this doubled the time needed.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class UltraCruciblePosition(CruciblePosition):\n",
    "    def moves(self, map: Map) -> t.Iterable[t.Self]:\n",
    "        orient = self.orientation\n",
    "        orthogonal = orient.orthogonal\n",
    "        new = type(self)\n",
    "\n",
    "        pos, hloss = self.pos, self.heat_loss\n",
    "        for step in range(1, 11):\n",
    "            pos += orient.negd\n",
    "            if (loss := map.get(pos)) is None:\n",
    "                break\n",
    "            hloss += loss\n",
    "            if step >= 4:\n",
    "                yield new(pos, orthogonal, hloss)\n",
    "\n",
    "        pos, hloss = self.pos, self.heat_loss\n",
    "        for step in range(1, 11):\n",
    "            pos += orient.posd\n",
    "            if (loss := map.get(pos)) is None:\n",
    "                break\n",
    "            hloss += loss\n",
    "            if step >= 4:\n",
    "                yield new(pos, orthogonal, hloss)\n",
    "\n",
    "\n",
    "assert test_map.least_heat_loss(UltraCruciblePosition) == 94"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Part 1: 1382\n"
     ]
    }
   ],
   "source": [
    "print(\"Part 1:\", map.least_heat_loss(UltraCruciblePosition))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}