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

# # Day 22, warped movements
# 
# We get to implement walking across a map, but this is a map with a twist: we need to identify edges that map to other edges.
# 
# We can't just map coordinates outside the traversable area to other coordinates; the direction of travel matters. E.g. in this section of the example, the coordinate at `X` is outside the map, and trying to 'walk' there from `A` puts you at `B`, but coming from `C` you need to move to `D`:
# 
# <pre><code>
#         ...#
#         .#..
#         #...
#        <span style="color: red">X</span><span style="color: steelblue">A</span>..<span style="color: dodgerblue">B</span>
# <span style="color: mediumpurple">G</span>..#...<span style="color: forestgreen">C</span>...#
# ........#...
# ..#....#....
# <span style="color: violet">E</span>......<span style="color: limegreen">D</span>..#<span style="color: mediumpurple">F</span>
# </code></pre>
# 
# As point `E` demonstrates, you can't use a simple map from point-on-the-edge to another point either, as, again, direction matters. Going left, you end up at `F`, going down you end up at `G`.
# 
# Finally, the example and the actual puzzle input differ in shape, and so I can only conclude that the actual layout will differ between different inputs as well. So we'll need a dynamic mapping solution!
# 
# So, what I did is track minima and maxima per column and row. Finding the next position is then a question of passing in the current position and direction, and we can then return the new position (if available).
# 

# In[1]:


# pyright: strict
import re
from contextlib import suppress
from enum import IntEnum
from typing import Final, NamedTuple, Self


class Pos(NamedTuple):
    x: int = 0
    y: int = 0

    def __add__(  # pyright: ignore[reportIncompatibleMethodOverride]
        self, other: Self
    ) -> "Pos":
        if not isinstance(other, Pos):
            return NotImplemented
        return Pos(self.x + other.x, self.y + other.y)

    def __neg__(self) -> "Pos":
        return Pos(-self.x, -self.y)

    def __mul__(  # pyright: ignore[reportIncompatibleMethodOverride]
        self, count: int
    ) -> "Pos":
        return Pos(self.x * count, self.y * count)


class Direction(IntEnum):
    L = -1
    R = 1


DIRECTIONS: Final[re.Pattern[str]] = re.compile(r"(L|R)")


class Facing(IntEnum):
    right = 0
    down = 1
    left = 2
    up = 3

    def rotate(self, dir: Direction) -> "Facing":
        return Facing((self.value + dir) % 4)

    def __invert__(self) -> "Facing":
        """Opposite direction"""
        return Facing((self.value + 2) % 4)

    def __radd__(  # pyright: ignore[reportIncompatibleMethodOverride]
        self, pos: Pos
    ) -> Pos:
        match self:
            case Facing.right:
                return Pos(pos.x + 1, pos.y)
            case Facing.down:
                return Pos(pos.x, pos.y + 1)
            case Facing.left:
                return Pos(pos.x - 1, pos.y)
            case Facing.up:
                return Pos(pos.x, pos.y - 1)


class BoardMap:
    _map: list[str]
    _shape: tuple[int, int]
    _start: Pos
    # indexed with facing and current y/x, gives new x/y coordinate when moving
    # in that direction. E.g. for pos x=17, y=42, _edges[l][42] -> right edge x,
    # _edges[u][17] -> bottom edge y
    _edges: tuple[list[int], list[int], list[int], list[int]]

    def __init__(self, map: str):
        self._map = lines = map.splitlines()
        self._shape = max(len(line) for line in lines), len(lines)
        self._start = self._init_edges()

    def _init_edges(self) -> Pos:
        lenx, leny = self._shape
        self._edges = edges = ([0] * leny, [leny] * lenx, [0] * leny, [0] * lenx)

        for y, row in enumerate(self._map):
            it = iter(enumerate(row))
            minx = next(x for x, char in it if char != " ")
            maxx = next((x for x, char in it if char == " "), len(row)) - 1
            edges[Facing.left][y] = maxx
            edges[Facing.right][y] = minx
            for x in range(minx, maxx + 1):
                edges[Facing.up][x] = max(edges[Facing.up][x], y)
                edges[Facing.down][x] = min(edges[Facing.down][x], y)
        return Pos(self._edges[Facing.right][0], 0)

    def move(self, pos: Pos, facing: Facing) -> tuple[Pos, Facing] | None:
        x, y = new = pos + facing
        map, edge, char = self._map, self._edges[facing], " "
        with suppress(IndexError):
            if x >= 0 <= y:
                char = map[y][x]
        # adjust for edges
        if char == " ":
            match facing:
                case Facing.up | Facing.down:
                    new = Pos(x, edge[x])
                case Facing.left | Facing.right:
                    new = Pos(edge[y], y)
            char = map[new.y][new.x]
        return None if char == "#" else (new, facing)

    def walk(self, steps: str) -> tuple[Pos, Facing]:
        facing, pos = Facing.right, self._start
        for value in DIRECTIONS.split(steps):
            match value:
                case "L" | "R":
                    facing = facing.rotate(Direction[value])
                case count:
                    for _ in range(int(count)):
                        if (next := self.move(pos, facing)) is None:
                            break
                        pos, facing = next
        return pos, facing

    def password(self, steps: str) -> int:
        pos, facing = self.walk(steps)
        row, col = pos.y + 1, pos.x + 1
        return facing + 1000 * row + 4 * col


example = """\
        ...#
        .#..
        #...
        ....
...#.......#
........#...
..#....#....
..........#.
        ...#....
        .....#..
        .#......
        ......#.

10R5L5R10L4R5L5
""".split("\n\n")
example_board = BoardMap(example[0])
assert example_board.password(example[1]) == 6032


# In[2]:


import aocd

bmap, steps = aocd.get_data(day=22, year=2022).split("\n\n")
board = BoardMap(bmap)
print("Part 1:", board.password(steps))


# ## Part 2, trouble cubed, not doubled
# 
# When wrapping, we now have to go round cube edges! That means we have to actually _rotate our orientation_ as well as go to a completely different edge.
# 
# What we've been given then, is the [net](https://en.wikipedia.org/wiki/Net_(polyhedron) of the cube. There are exactly 11 possible different arrengements of the 6 sides (not counting rotations and reflections), but we don't actually need to know which one our puzzle input uses. Instead, you can treat the faces and the vertices (the lines that form the edge between two faces) form a _graph_, and we can just assign one of the faces to the first square you come across in the puzzle, and assign vertices to its 4 edges. From there you can then determine where the other faces are in the flat map and how their edges relate.
# 
# Once you have that information, it's a question of basic linear algebra: matrix multiplications with [transformation matrices](https://en.wikipedia.org/wiki/Transformation_matrix) to transform points on one edge to the points on the corresponding edge on another square. We can pre-calculate the required transformation matrix for any given edge mapping, together with a rotation you need to apply to the facing each time you cross an edge.
# 
# I had to refactor the solution for part 1 only slightly to accomodate the part 2 subclass; the `move()` method now returns `(Pos, Facing)` for each step that's possible, not just `Pos`.
# 

# In[3]:


from collections import deque
from itertools import product
from math import sqrt
from typing import Literal, TypeAlias

# Cube graph; 6 faces with 4 vertices each, vertices shared between faces.
# 6 faces numbered 0 - 5, arranged so opposite faces sum to 5.
# 12 vertices labelled a - n (skipping i, l and o for readability), in clockwise
# order starting at the Facing.right position.
Vertex: TypeAlias = Literal["a", "b", "c", "d", "e", "f", "g", "h", "j", "k", "m", "n"]
Face: TypeAlias = int
VertexMap: TypeAlias = dict[Vertex, Face]


CUBE_FACES: Final[
    tuple[VertexMap, VertexMap, VertexMap, VertexMap, VertexMap, VertexMap]
] = (
    {"a": 1, "b": 2, "c": 4, "d": 3},
    {"e": 3, "f": 5, "g": 2, "a": 0},
    {"g": 1, "h": 5, "j": 4, "b": 0},
    {"m": 4, "n": 5, "e": 1, "d": 0},
    {"j": 2, "k": 5, "m": 3, "c": 0},
    {"k": 4, "h": 2, "f": 1, "n": 3},
)


class Transformation(NamedTuple):
    xs: tuple[int, int, int]
    ys: tuple[int, int, int]
    zs: tuple[int, int, int]

    def __matmul__(self, rhs: Self) -> Self:
        if not isinstance(rhs, Transformation):
            return NotImplemented
        cols = list(zip(*rhs))
        return type(self)(
            tuple(sum(x * c for x, c in zip(self.xs, col)) for col in cols),
            tuple(sum(y * c for y, c in zip(self.ys, col)) for col in cols),
            tuple(sum(z * c for z, c in zip(self.zs, col)) for col in cols),
        )

    def __rmatmul__(self, pos: Pos) -> Pos:
        """Transform a point.

        Technically, the operation is T @ P, but P @ T is more convenient here
        as it keeps the implementation simpler.

        """
        if not isinstance(pos, Pos):
            return NotImplemented
        v = (*pos, 1)
        return Pos(
            sum(x * c for x, c in zip(self.xs, v)),
            sum(y * c for y, c in zip(self.ys, v)),
        )

    @classmethod
    def rotate(cls, dir: Direction) -> Self:
        match dir:
            case Direction.L:
                return cls((0, 1, 0), (-1, 0, 0), (0, 0, 1))
            case Direction.R:
                return cls((0, -1, 0), (1, 0, 0), (0, 0, 1))

    @classmethod
    def translate(cls, x: int, y: int) -> Self:
        return cls((1, 0, x), (0, 1, y), (0, 0, 1))


RotR = Transformation.rotate(Direction.R)
RotL = Transformation.rotate(Direction.L)
Rot180 = RotR @ RotR


class CubeBoardMap(BoardMap):
    _cube_len: int
    # current facing and position, moving means you end up at new facing
    # and transformation matrix gives new position.
    _vertex_map: dict[tuple[Facing, Pos], tuple[Facing, Transformation]]

    def _init_edges(self):
        points = sum(len(section) for line in self._map for section in line.split())
        # vertex length; side of one of the 6 cubes
        vl = self._cube_len = int(sqrt(points // 6))
        cw, ch = self._shape[1] // vl, self._shape[0] // vl
        # traverse the cube net graph to locate all matched and unmatched vertices
        pos = next(
            Pos(cx, cy)
            for cy, cx in product(range(ch), range(cw))
            if self._map[cy * vl][cx * vl] != " "
        )
        positioned: dict[Face, tuple[Pos, dict[Vertex, Facing]]] = {
            0: (pos, dict(zip(CUBE_FACES[0], Facing)))
        }
        disconnected: dict[Face, set[Vertex]] = {f: set() for f in range(6)}
        queue: deque[Face] = deque([0])
        seen: set[Face] = {0}
        while queue:
            face = queue.popleft()
            pos, vmap = positioned[face]
            current_vertices = CUBE_FACES[face]
            dis = disconnected[face]
            for vertex, vdir in vmap.items():
                opos = pos + vdir
                char = " "
                with suppress(IndexError):
                    if opos.x >= 0 <= opos.y:
                        char = self._map[opos.y * vl][opos.x * vl]
                if char == " ":
                    dis.add(vertex)
                    continue
                # Cube face at this position
                face = current_vertices[vertex]
                if face in seen:
                    continue
                seen.add(face)
                vertices = tuple(CUBE_FACES[face])
                # rotate to match the correct orientation
                delta = (vertices.index(vertex) - ~vdir) % 4
                vertices = vertices[delta:] + vertices[:delta]
                positioned[face] = (opos, dict(zip(vertices, Facing)))
                queue.append(face)

        # build disconnected vertex edge transformations
        self._vertex_map = vmap = {}
        for face, vertexes in disconnected.items():
            for vertex in vertexes:
                other = CUBE_FACES[face][vertex]
                hpos, hvdir = positioned[face]
                opos, ovdir = positioned[other]
                hdir, odir = hvdir[vertex], ovdir[vertex]
                # adjust positions to represent the first corner in Facing order
                # e.g. Facing.right -> top-right corner, Facing.down ->
                # bottom-right, etc.
                hcorner = (hpos * vl) + Pos(
                    vl - 1 if hdir in {Facing.right, Facing.down} else 0,
                    vl - 1 if hdir in {Facing.down, Facing.left} else 0,
                )
                # Use **opposite** corner of the vertex as this is the direction
                # to _enter_ the cube.
                ocorner = (opos * vl) + Pos(
                    vl - 1 if odir in {Facing.up, Facing.right} else 0,
                    vl - 1 if odir in {Facing.right, Facing.down} else 0,
                )
                trans = Transformation.translate(*-hcorner)
                match (hdir - ~odir) % 4:
                    case 1:
                        trans = RotL @ trans
                    case 2:
                        trans = Rot180 @ trans
                    case 3:
                        trans = RotR @ trans
                trans = Transformation.translate(*ocorner) @ trans
                match hdir:
                    case Facing.right:
                        points = (hcorner + Pos(0, y) for y in range(vl))
                    case Facing.down:
                        points = (hcorner + Pos(-x, 0) for x in range(vl))
                    case Facing.left:
                        points = (hcorner + Pos(0, -y) for y in range(vl))
                    case Facing.up:
                        points = (hcorner + Pos(x, 0) for x in range(vl))
                for point in points:
                    vmap[hdir, point] = (~odir, trans)
        return positioned[0][0] * vl

    def move(self, pos: Pos, facing: Facing) -> tuple[Pos, Facing] | None:
        if match := self._vertex_map.get((facing, pos)):
            facing, trans = match
            x, y = new = pos @ trans
        else:
            x, y = new = pos + facing
        # x and y are guaranteed to be a valid cube face in the map
        return None if self._map[y][x] == "#" else (new, facing)


example_cube = CubeBoardMap(example[0])
assert example_cube.password(example[1]) == 5031


# In[4]:


cube = CubeBoardMap(bmap)
print("Part 2:", cube.password(steps))