Metro map

"""An invented metro map in drafted octilinear lines."""

from collections.abc import Sequence
from dataclasses import dataclass
from itertools import pairwise

from walldye import (
    ACCENT,
    BG,
    BG_ALT,
    BG_DEEP,
    MUTED,
    UI,
    UI_ALT,
    UI_HI,
    Canvas,
    Color,
    P,
    Path,
    Vec,
    by_regime,
    design,
    mix,
)
from walldye.geom import Affine

G = 60  # grid step
LANE = 14  # spacing of parallel lines sharing a corridor
LW = 8
R = 30  # bend radius
TICK, BAR = 13, 12  # station tick length; half-length of a terminus bar
HUB, HUB_CORE = 20, 14  # interchange capsule width and its hollow

type Node = tuple[int, int]

# (paint, width, grid polyline); order also sets lane order within shared corridors
LINES: tuple[tuple[Color, float, tuple[Node, ...]], ...] = (
    (UI, LW, ((10, 5), (18, 5), (20, 7), (20, 11), (18, 13), (10, 13), (8, 11), (8, 7), (10, 5))),
    (UI_ALT, LW, ((5, 9), (17, 9), (19, 11), (26, 11))),
    (ACCENT, 10, ((5, 15), (8, 15), (14, 9), (17, 9), (23, 3), (26, 3))),
    (UI, LW, ((14, 2), (14, 16))),
    (UI_ALT, LW, ((6, 3), (8, 3), (10, 5), (18, 5), (21, 2))),
    (UI, LW, ((11, 16), (11, 14), (10, 13), (8, 13), (6, 11), (6, 6))),
    (UI, LW, ((16, 16), (18, 14), (18, 13), (23, 13), (25, 15))),
)
PICKED = 2  # the accent line; its interchanges get the brighter rim
# river bends in grid units; its ends run on past the canvas edges
RIVER = ((3, 12), (5, 10), (9.5, 10), (11.5, 12), (19.5, 12), (21.5, 14), (23.5, 14), (25, 12.5))
MID = Vec(15.5 * G, 9 * G)  # center of the network's grid box
# a recess below the ground; on paper BG_DEEP is paler than the page, so shade toward the ink
WATER = by_regime(BG_DEEP, mix(BG, BG_ALT, 0.6))


@dataclass
class Step:
    """One grid step of a line: its lane offset at each end, how many lines share the corridor,
    and the corridor's unit normal, which is the same whichever way a line travels it."""

    a: Node
    b: Node
    o0: float
    o1: float
    shared: int
    nrm: Vec

    @property
    def heading(self) -> Node:
        return (self.b[0] - self.a[0], self.b[1] - self.a[1])


# half a step: the grid node it starts at, if any, and its two ends
type Half = tuple[Node | None, Vec, Vec]


def edges(nodes: Sequence[Node]) -> list[tuple[Node, Node]]:
    """Split a grid polyline into unit steps."""
    out: list[tuple[Node, Node]] = []
    for (x0, y0), (x1, y1) in pairwise(nodes):
        n = max(abs(x1 - x0), abs(y1 - y0))
        dx, dy = (x1 - x0) // n, (y1 - y0) // n
        out += [
            ((x0 + k * dx, y0 + k * dy), (x0 + (k + 1) * dx, y0 + (k + 1) * dy)) for k in range(n)
        ]
    return out


def corridor(a: Node, b: Node) -> tuple[tuple[Node, Node], Vec]:
    """The undirected key of a step and its unit normal."""
    if b < a:
        a, b = b, a
    return (a, b), Vec(b[0] - a[0], b[1] - a[1]).perp().unit()


def cross(u: Vec, v: Vec) -> float:
    return u.x * v.y - u.y * v.x


def rounded(pts: Sequence[Vec], r: float) -> Path:
    """A polyline whose corners are cut by quadratic curves starting up to `r` back along each
    leg (at most half of it)."""
    d = P().M(pts[0])
    for p0, p1, p2 in zip(pts, pts[1:], pts[2:]):
        l0, l1 = abs(p0 - p1), abs(p2 - p1)
        a = p1 + (p0 - p1) * min(r, l0 / 2) / l0
        b = p1 + (p2 - p1) * min(r, l1 / 2) / l1
        d.L(a).Q(p1, b)
    return d.L(pts[-1])


def route(
    k: int, nodes: tuple[Node, ...], lanes: dict[tuple[Node, Node], list[int]]
) -> tuple[list[Vec], dict[Node, tuple[Vec, ...]], list[Step]]:
    """Line `k` in pixels: its corner polyline, its lane points at each grid node it passes, and
    its unit steps. `lanes` lists the lines using each corridor, in lane order."""
    closed = nodes[0] == nodes[-1]
    steps: list[Step] = []
    for a, b in edges(nodes):
        key, nrm = corridor(a, b)
        users = lanes[key]
        off = (users.index(k) - (len(users) - 1) / 2) * LANE
        steps.append(Step(a, b, off, off, len(users), nrm))
    # a lane change on a straight happens mid-step on the less shared side, never inside an interchange
    n = len(steps)
    for i in range(n if closed else n - 1):
        s0, s1 = steps[i], steps[(i + 1) % n]
        if s0.heading == s1.heading and s0.o1 != s1.o0:
            if s0.shared < s1.shared:
                s0.o1 = s1.o0
            else:
                s1.o0 = s0.o1

    def px(p: Vec, nrm: Vec, off: float) -> Vec:
        return p * G + nrm * off

    halves: list[Half] = []
    for st in steps:
        a, b = Vec(*st.a), Vec(*st.b)
        m = (a + b) / 2
        halves += [
            (st.a, px(a, st.nrm, st.o0), px(m, st.nrm, st.o0)),
            (None, px(m, st.nrm, st.o1), px(b, st.nrm, st.o1)),
        ]
    poly: list[Vec] = []
    at: dict[Node, tuple[Vec, ...]] = {}
    for (_, a0, b0), (node, a1, b1) in zip(
        halves, halves[1:] + halves[:1] if closed else halves[1:]
    ):
        d0, d1 = b0 - a0, b1 - a1
        c = cross(d0, d1)
        if node is not None:
            at[node] = (b0, a1)
        if abs(c) < 1e-9:
            if abs(a1 - b0) > 0.5:  # lane change on a straight: a 45-degree jog
                j = d0.unit() * (abs(a1 - b0) / 2)
                poly += [b0 - j, a1 + j]
            continue
        poly.append(a0 + d0 * (cross(a1 - a0, d1) / c))
    if closed:  # start mid-segment so every corner gets rounded
        m = (poly[-1] + poly[0]) / 2
        return [m, *poly, m], at, steps
    at[nodes[0]], at[nodes[-1]] = (halves[0][1],), (halves[-1][2],)
    return [halves[0][1], *poly, halves[-1][2]], at, steps


@design(aspects="any")
def draw(s: Canvas) -> None:
    c = s.pick(landscape=(31 / 64, 0.5), portrait=(0.5, 0.5), snap=30)
    if s.landscape:
        place = Affine.translate(c.x - MID.x, c.y - MID.y)
    else:  # mirrored across the diagonal, so the river runs down the screen
        place = Affine(0, 1, 1, 0, c.x - MID.y, c.y - MID.x)
    back = place.inverse()
    x0, x1 = back((0, 0)).x - G, back((s.w, s.h)).x + G  # past both canvas edges

    lanes: dict[tuple[Node, Node], list[int]] = {}
    for k, (*_, nodes) in enumerate(LINES):
        for a, b in edges(nodes):
            lanes.setdefault(corridor(a, b)[0], []).append(k)
    routes = [route(k, nodes, lanes) for k, (*_, nodes) in enumerate(LINES)]

    # an interchange wherever lines meet at a node without all running straight through side by side
    meets: dict[Node, dict[int, tuple[Vec, ...]]] = {}
    dirs: dict[Node, dict[int, set[Node]]] = {}
    for k, ((*_, nodes), (_, at, _)) in enumerate(zip(LINES, routes)):
        for node, ps in at.items():
            meets.setdefault(node, {})[k] = ps
        for a, b in edges(nodes):
            dirs.setdefault(a, {}).setdefault(k, set()).add((b[0] - a[0], b[1] - a[1]))
            dirs.setdefault(b, {}).setdefault(k, set()).add((a[0] - b[0], a[1] - b[1]))
    hubs = {
        n: ps
        for n, ps in meets.items()
        if len(ps) > 1 and len({frozenset(v) for v in dirs[n].values()}) > 1
    }

    with s.group(transform=place):
        river = [Vec(x0, RIVER[0][1] * G), *(Vec(x, y) * G for x, y in RIVER)]
        s.stroke(rounded([*river, Vec(x1, RIVER[-1][1] * G)], 40), WATER, 40, join="round")

        rnd = s.rng(4)
        for k, ((paint, width, nodes), (poly, _, steps)) in enumerate(zip(LINES, routes)):
            ends = P()
            if nodes[0] != nodes[-1]:
                for end, nrm in ((poly[0], steps[0].nrm), (poly[-1], steps[-1].nrm)):
                    ends.M(end - nrm * BAR).L(end + nrm * BAR)
            # station ticks mid-step, on the outer side of any corridor, clear of bends, jogs
            # and interchanges
            ticks = P()
            for n, st in enumerate(steps):
                if (
                    st.a in hubs
                    or st.b in hubs
                    or st.o0 != st.o1
                    or n in (0, len(steps) - 1)
                    or rnd.random() < 0.45
                ):
                    continue
                side = 1 if st.o0 > 0 else -1 if st.o0 < 0 else (1 if k % 2 else -1)
                m = (Vec(*st.a) + st.b) / 2 * G + st.nrm * st.o0
                ticks.M(m).L(m + st.nrm * (side * TICK))
            s.stroke(rounded(poly, R), paint, width, cap="round", join="round")
            s.stroke(ends, paint, 5, cap="round")
            s.stroke(ticks, paint, 4)

        # capsules spanning every lane through the node; the accent line's get the brighter rim
        cores = P()
        with s.buckets((UI_HI, MUTED), "stroke", stroke_width=HUB, stroke_linecap="round") as rims:
            for ps in hubs.values():
                pl = [p for v in ps.values() for p in v]
                a, b = max(((p, q) for p in pl for q in pl), key=lambda pq: abs(pq[1] - pq[0]))
                if abs(b - a) <= 0.5:  # one lane: nudge the end so the round caps draw a dot
                    b = a + (0.01, 0)
                rims[1 if PICKED in ps else 0].M(a).L(b)
                cores.M(a).L(b)
        s.stroke(cores, BG, HUB_CORE, cap="round")