engine/traffic.py

"""NEMOCITY traffic — the Python static-assignment mirror.

Same trip generation + A* + demand semantics as web/city/traffic.js, from the
same event list and constants, so the ENGINE owns the facts the LLM sees (the
live car queueing client-side is presentation and may drift a car or two).

Determinism contract with the JS mirror:
* mulberry32 (bit-exact port of web/city/rng.js); ONE draw per commuter,
  seeded `building.seed ^ i`;
* trip order = building event order, commuter index ascending;
* destination roulette over completed jobs/attract buildings in event order,
  weight = (jobs + 2*attract) / (1 + manhattanCentroidDist^1.5);
* door = adjacent road cell with smallest (cz, cx);
* incremental loading: each commuter routes by A* with
  cost(cell) = (1/speed) * (1 + 2*demandRatio(cell)) over demand loaded so
  far, then adds DEMAND_PER_COMMUTER to every cell of its path.

Stats are evaluated AT PEAK RUSH (rushFactor = 1.0). Pure stdlib.
"""

from __future__ import annotations

import heapq
import time
from dataclasses import dataclass, field
from typing import Any, Optional

from . import constants as C
from .city import WATER_CELLS, Building, Cell, CityState, in_bounds, neighbors4
from .tools import street_name_for

_MIN_CELL_COST = 1.0 / C.ROAD_CLASSES["avenue"]["speed"]  # admissible A* heuristic


def mulberry32(seed: int):
    """Bit-exact port of web/city/rng.js mulberry32 — keep byte-stable."""
    a = seed & 0xFFFFFFFF

    def rng() -> float:
        nonlocal a
        a = (a + 0x6D2B79F5) & 0xFFFFFFFF
        t = a
        t = ((t ^ (t >> 15)) * (t | 1)) & 0xFFFFFFFF
        t = (t ^ ((t + (((t ^ (t >> 7)) * (t | 61)) & 0xFFFFFFFF)) & 0xFFFFFFFF)) & 0xFFFFFFFF
        return ((t ^ (t >> 14)) & 0xFFFFFFFF) / 4294967296.0

    return rng


# -------------------------------------------------------------------- road net

@dataclass
class RoadNet:
    """Road cells + class/name, optionally with a candidate fix applied."""

    cells: dict[Cell, tuple[str, str]]  # cell -> (klass, name)

    @classmethod
    def from_city(cls, city: CityState, fix: Optional[dict] = None) -> "RoadNet":
        cells = {cell: (rc.klass, rc.name) for cell, rc in city.roads.items()}
        if fix:
            if fix["action"] == "new_road":
                for cx, cz in fix["cells"]:
                    cells[(cx, cz)] = (fix["klass"], fix["name"])
            elif fix["action"] == "upgrade_avenue":
                for cx, cz in fix["cells"]:
                    cell = (cx, cz)
                    if cell in cells:
                        cells[cell] = ("avenue", cells[cell][1])
        return cls(cells=cells)

    def capacity(self, cell: Cell) -> float:
        return C.ROAD_CLASSES[self.cells[cell][0]]["capacity"]

    def speed(self, cell: Cell) -> float:
        return C.ROAD_CLASSES[self.cells[cell][0]]["speed"]

    def name(self, cell: Cell) -> str:
        return self.cells[cell][1]

    def door(self, b: Building) -> Optional[Cell]:
        adjacent = {
            n for cell in b.cells for n in neighbors4(cell) if n in self.cells
        }
        if not adjacent:
            return None
        return min(adjacent, key=lambda c: (c[1], c[0]))


# -------------------------------------------------------------------- trip gen

def trip_table(city: CityState, now_s: float) -> list[tuple[Building, Building]]:
    """Static OD pairs (gravity model, seeded roulette). Completed buildings
    only on both ends; recomputed only when buildings change."""
    completed = [b for b in city.buildings if b.progress(now_s) >= 1.0]
    homes = [b for b in completed if C.BUILDINGS[b.kind]["residents"] > 0]
    dests = [
        b for b in completed
        if C.BUILDINGS[b.kind]["jobs"] + C.BUILDINGS[b.kind]["attract"] > 0
    ]
    if not dests:
        return []
    trips: list[tuple[Building, Building]] = []
    for home in homes:
        hx, hz = home.centroid
        weights: list[float] = []
        for dest in dests:
            dx, dz = dest.centroid
            dist = abs(dx - hx) + abs(dz - hz)
            spec = C.BUILDINGS[dest.kind]
            weights.append(
                (spec["jobs"] + C.TRIP_ATTRACT_FACTOR * spec["attract"])
                / (1 + dist ** C.TRIP_DIST_EXP)
            )
        total = sum(weights)
        if total <= 0:
            continue
        for i in range(C.BUILDINGS[home.kind]["residents"]):
            r = mulberry32((home.seed ^ i) & 0xFFFFFFFF)() * total
            acc = 0.0
            chosen = dests[-1]
            for dest, weight in zip(dests, weights):
                acc += weight
                if r < acc:
                    chosen = dest
                    break
            trips.append((home, chosen))
    return trips


# ------------------------------------------------------------------ assignment

@dataclass
class Assignment:
    net: RoadNet
    demand: dict[Cell, float] = field(default_factory=dict)
    paths: list[tuple[Cell, Cell, tuple[Cell, ...]]] = field(default_factory=list)
    traffic_index: int = 0

    def ratio(self, cell: Cell) -> float:
        return self.demand.get(cell, 0.0) / self.net.capacity(cell)

    @property
    def max_ratio(self) -> float:
        if not self.net.cells:
            return 0.0
        return max((self.ratio(c) for c in self.net.cells), default=0.0)

    def top_cells(self, n: int) -> list[Cell]:
        return sorted(
            self.net.cells,
            key=lambda c: (-self.ratio(c), c[1], c[0]),
        )[:n]


def _astar(net: RoadNet, demand: dict[Cell, float], start: Cell, goal: Cell
           ) -> Optional[list[Cell]]:
    if start == goal:
        return [start]

    def cell_cost(cell: Cell) -> float:
        ratio = demand.get(cell, 0.0) / net.capacity(cell)
        return (1.0 / net.speed(cell)) * (1 + C.CONGESTION_COST_FACTOR * ratio)

    def h(cell: Cell) -> float:
        return (abs(cell[0] - goal[0]) + abs(cell[1] - goal[1])) * _MIN_CELL_COST

    g: dict[Cell, float] = {start: 0.0}
    prev: dict[Cell, Cell] = {}
    heap: list[tuple[float, int, int, Cell]] = [(h(start), start[1], start[0], start)]
    closed: set[Cell] = set()
    while heap:
        _, _, _, cell = heapq.heappop(heap)
        if cell in closed:
            continue
        if cell == goal:
            path = [cell]
            while path[-1] != start:
                path.append(prev[path[-1]])
            path.reverse()
            return path
        closed.add(cell)
        for n in neighbors4(cell):
            if n not in net.cells or n in closed:
                continue
            cand = g[cell] + cell_cost(n)
            if cand < g.get(n, float("inf")):
                g[n] = cand
                prev[n] = cell
                heapq.heappush(heap, (cand + h(n), n[1], n[0], n))
    return None


def assign(city: CityState, now_s: float, fix: Optional[dict] = None) -> Assignment:
    """Incremental static assignment at peak rush. <2s for a genesis city."""
    net = RoadNet.from_city(city, fix)
    out = Assignment(net=net)
    for home, dest in trip_table(city, now_s):
        od, dd = net.door(home), net.door(dest)
        if od is None or dd is None:
            continue
        path = _astar(net, out.demand, od, dd)
        if path is None:
            continue
        for cell in path:
            out.demand[cell] = out.demand.get(cell, 0.0) + C.DEMAND_PER_COMMUTER
        out.paths.append((od, dd, tuple(path)))
    ratios = sorted((out.ratio(c) for c in net.cells), reverse=True)
    top = ratios[: C.TRAFFIC_TOP_CELLS]
    if top:
        # pad: a city with fewer road cells than the window still reads honestly
        mean = sum(top) / C.TRAFFIC_TOP_CELLS
        out.traffic_index = int(round(100 * mean))
    return out


# ----------------------------------------------------------------------- stats

def _intersection_name(net: RoadNet, cell: Cell) -> str:
    names = [net.name(cell)]
    for n in neighbors4(cell):
        if n in net.cells and net.name(n) not in names:
            names.append(net.name(n))
    return " & ".join(names[:2])


def stats(city: CityState, assignment: Assignment) -> dict:
    net = assignment.net
    top = [
        c for c in assignment.top_cells(5) if assignment.ratio(c) > 0
    ]
    worst = top[0] if top else None
    intersections = [
        c for c in net.cells
        if sum(1 for n in neighbors4(c) if n in net.cells) >= 3
    ]
    worst_x = max(
        intersections, key=lambda c: (assignment.ratio(c), -c[1], -c[0]), default=None
    )
    return {
        "traffic_index": assignment.traffic_index,
        "max_ratio": round(assignment.max_ratio, 2),
        "worst": net.name(worst) if worst else None,
        "worst_intersection": _intersection_name(net, worst_x) if worst_x else None,
        "top": [
            {
                "street": net.name(c),
                "cell": [c[0], c[1]],
                "ratio": round(assignment.ratio(c), 2),
                "klass": net.cells[c][0],
            }
            for c in top
        ],
    }


# ------------------------------------------------------------------ candidates

def _bypass_candidate(
    city: CityState, assignment: Assignment, now_s: float
) -> Optional[dict]:
    """Engine-routed relief road around the worst cell: Dijkstra over the grid,
    new cells cost extra, water at 4x (a fix may deliberately bridge), jammed
    cells forbidden."""
    net = assignment.net
    jammed = {
        c for c in net.cells if assignment.ratio(c) >= C.FIX_AVOID_RATIO
    }
    if not jammed:
        return None
    worst = min(jammed, key=lambda c: (-assignment.ratio(c), c[1], c[0]))
    od_counts: dict[tuple[Cell, Cell], int] = {}
    for od, dd, path in assignment.paths:
        if worst in path:
            od_counts[(od, dd)] = od_counts.get((od, dd), 0) + 1
    if not od_counts:
        return None
    (start, goal), _ = min(
        od_counts.items(), key=lambda kv: (-kv[1], kv[0][0][1], kv[0][0][0])
    )

    def cost(cell: Cell) -> Optional[float]:
        if cell in jammed and cell not in (start, goal):
            return None
        if cell in net.cells:
            return 1.0 / net.speed(cell)
        if cell in city.building_cells:
            return None
        if cell in WATER_CELLS:
            return C.FIX_NEW_CELL_COST * C.FIX_WATER_COST_MULT
        return C.FIX_NEW_CELL_COST

    dist: dict[Cell, float] = {start: 0.0}
    prev: dict[Cell, Cell] = {}
    heap: list[tuple[float, int, int, Cell]] = [(0.0, start[1], start[0], start)]
    closed: set[Cell] = set()
    while heap:
        d, _, _, cell = heapq.heappop(heap)
        if cell in closed:
            continue
        if cell == goal:
            break
        closed.add(cell)
        for n in neighbors4(cell):
            if n in closed:
                continue
            c = cost(n)
            if c is None:
                continue
            cand = d + c
            if cand < dist.get(n, float("inf")):
                dist[n] = cand
                prev[n] = cell
                heapq.heappush(heap, (cand, n[1], n[0], n))
    if goal not in prev and goal != start:
        return None
    path = [goal]
    while path[-1] != start:
        path.append(prev[path[-1]])
    path.reverse()
    new_cells = [c for c in path if c not in net.cells]
    if not new_cells or len(new_cells) > C.FIX_MAX_NEW_CELLS:
        return None
    name = street_name_for(f"fix:{city.road_version}:bypass")
    bridges = sum(1 for c in new_cells if c in WATER_CELLS)
    fix = {
        "action": "new_road",
        "cells": [[c[0], c[1]] for c in new_cells],
        "klass": "avenue",
        "name": name,
    }
    desc = f"new avenue {name}, {len(new_cells)} cells"
    if bridges:
        desc += f" incl. a {bridges}-cell second bridge"
    fix["desc"] = desc
    fix["predicted_index"] = assign(city, now_s, fix).traffic_index
    return fix


def _upgrade_candidates(
    city: CityState, assignment: Assignment, now_s: float, limit: int = 2
) -> list[dict]:
    net = assignment.net
    names: list[str] = []
    for cell in assignment.top_cells(C.TRAFFIC_TOP_CELLS):
        klass, name = net.cells[cell]
        if klass == "street" and assignment.ratio(cell) > 0 and name not in names:
            names.append(name)
        if len(names) >= limit:
            break
    out = []
    for name in names:
        cells = [
            [c[0], c[1]]
            for c, (klass, n) in net.cells.items()
            if n == name and klass == "street"
        ]
        if not cells:
            continue
        cells.sort(key=lambda c: (c[1], c[0]))
        fix = {
            "action": "upgrade_avenue",
            "cells": cells,
            "klass": "avenue",
            "name": name,
            "desc": f"upgrade {name} to a 6-lane avenue ({len(cells)} cells)",
        }
        fix["predicted_index"] = assign(city, now_s, fix).traffic_index
        out.append(fix)
    return out


def candidates(city: CityState, assignment: Assignment, now_s: float) -> list[dict]:
    """2-4 pre-validated fixes, best predicted first, ids F1..Fn."""
    found: list[dict] = []
    bypass = _bypass_candidate(city, assignment, now_s)
    if bypass:
        found.append(bypass)
    found.extend(_upgrade_candidates(city, assignment, now_s))
    found.sort(key=lambda f: (f["predicted_index"], f["action"], f["name"]))
    for i, fix in enumerate(found[:4]):
        fix["id"] = f"F{i + 1}"
    return found[:4]


def snapshot(city: CityState, now_s: Optional[float] = None) -> tuple[dict, list[dict]]:
    """(stats, candidates) for the fix mechanic — the engine owns the facts."""
    if now_s is None:
        now_s = time.time()
    assignment = assign(city, now_s)
    st = stats(city, assignment)
    cands = candidates(city, assignment, now_s) if assignment.max_ratio > 0 else []
    return st, cands