Where we are: The experiment record

"""AGENT-EDITABLE — model architecture + deployment inference wrapper.

The autoresearch loop may rewrite anything here (architecture, loss,
confidence mechanism, unified vs. dispatcher+specialists, pretrained init
per CLAUDE.md §3/§9) as long as:
  - train.py exports one ONNX per area at <out>/models/<area>.onnx taking a
    1x3x128x128 float32 input in [0,1] and returning [[u, v, conf]] with
    u, v normalized map coords and conf in [0,1];
  - the export passes pipeline/score.py's frozen deployment gates.

Current: localization as MAP-CELL CLASSIFICATION, from-scratch init, decoded by
LOCAL BELIEF POOLING. The map is cut into CELL_PX-square cells (Berlin:
55x54 = 2970); a conv trunk feeds a C-way softmax over cells plus a within-cell
offset head. The deployed decode pools belief over the 3x3 block of cells
centred on the argmax cell: confidence is the softmax mass on that whole
neighborhood, and the fix is the mass-weighted centroid of those cell centres,
nudged by the offset. Eval views stand 11-17 m off-lattice, so a view near a
cell border honestly splits its mass between adjacent cells — that is label
ambiguity, not place ambiguity, and pooling makes confidence measure "does the
model know WHERE it is" rather than "which arbitrary tile label applies".
Belief scattered across DISTANT cells still falls below threshold, so genuine
far-field ambiguity remains a safe abstention rather than a false fix. Cell
size is chosen so the parameterization alone suffices: a 128 m cell's
half-diagonal is 90.5 m, inside the 100 m usable radius, so naming the right
cell is already a usable fix. The TRAINING TARGETS are decode-consistent: the
cell target is a bilinear tent distribution over the <=4 nearest cell centres
(so a border view is taught the honest split its pooled decode will produce,
not an arbitrary one-cell label), and the offset target is the residual the
pooled centroid leaves rather than the position that centroid already encodes.

The whole decode (softmax -> argmax -> neighborhood mask -> pooled centroid ->
offset -> conf) lives INSIDE the exported graph, because the scorer runs the
ONNX directly.
"""

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

CELL_PX = 128  # map cell edge in raster px (== metres at this GSD)


def grid_dims(meta: dict) -> tuple[int, int]:
    """(gw, gh): map-cell grid covering the area's raster."""
    return (math.ceil(meta["width"] / CELL_PX), math.ceil(meta["height"] / CELL_PX))


class CellLocNet(nn.Module):
    """Conv trunk -> GAP -> (C-way cell logits, within-cell offset).

    forward() is the DEPLOYED path and returns [B, 3] = (u, v, conf).
    training_outputs() exposes the raw (logits, offset) for the loss.
    """

    def __init__(self, meta: dict):
        super().__init__()
        gw, gh = grid_dims(meta)
        self.gw, self.gh = gw, gh
        self.n_cells = gw * gh
        width, height = meta["width"], meta["height"]

        chans = [3, 32, 64, 128, 160]
        layers = []
        for cin, cout in zip(chans, chans[1:]):
            layers += [nn.Conv2d(cin, cout, 3, stride=2, padding=1),
                       nn.BatchNorm2d(cout), nn.ReLU(inplace=True)]
        self.features = nn.Sequential(*layers)
        self.cls_head = nn.Linear(chans[-1], self.n_cells)
        self.off_head = nn.Linear(chans[-1], 2)

        # Cell-centre lookup tables, in normalized map coords. Registered as
        # buffers so they bake into the ONNX graph as constants.
        idx = np.arange(self.n_cells)
        cell_u = ((idx % gw) + 0.5) * CELL_PX / width
        cell_v = ((idx // gw) + 0.5) * CELL_PX / height
        self.register_buffer("cell_u", torch.tensor(cell_u, dtype=torch.float32))
        self.register_buffer("cell_v", torch.tensor(cell_v, dtype=torch.float32))
        self.scale_u = CELL_PX / width
        self.scale_v = CELL_PX / height

    def training_outputs(self, x):
        f = self.features(x).mean(dim=(2, 3))
        # tanh*0.5 keeps offsets inside the cell, in (-0.5, 0.5) cell units.
        return self.cls_head(f), torch.tanh(self.off_head(f)) * 0.5

    def forward(self, x):
        # Local belief pooling: confidence is the softmax mass on the 3x3 block
        # of cells centred on the argmax cell, and the fix is that block's
        # mass-weighted centroid (plus the within-cell offset). A border view
        # splitting mass between adjacent cells still names one place, so it
        # stays a fix; mass scattered across distant cells stays an abstention.
        logits, offset = self.training_outputs(x)
        p = F.softmax(logits, dim=1)
        idx = p.argmax(dim=1)
        win_u = self.cell_u.index_select(0, idx)
        win_v = self.cell_v.index_select(0, idx)
        # 1.6 cell-pitches per axis (Chebyshev) selects exactly the 3x3 block
        # with float-safety margin — the next ring sits at 2.0 pitches. Edge
        # cells simply have fewer neighbours in range; nothing is counted twice.
        du = (self.cell_u.unsqueeze(0) - win_u.unsqueeze(1)).abs()
        dv = (self.cell_v.unsqueeze(0) - win_v.unsqueeze(1)).abs()
        mask = ((du <= 1.6 * self.scale_u) & (dv <= 1.6 * self.scale_v)).float()
        q = p * mask
        conf = q.sum(dim=1)  # sub-sum of a softmax: in [0,1], >= single-cell max
        w = q / conf.unsqueeze(1).clamp_min(1e-9)
        u = (w * self.cell_u.unsqueeze(0)).sum(dim=1) + offset[:, 0] * self.scale_u
        v = (w * self.cell_v.unsqueeze(0)).sum(dim=1) + offset[:, 1] * self.scale_v
        return torch.stack([u.clamp(0.0, 1.0), v.clamp(0.0, 1.0), conf], dim=1)


def build_model(meta: dict) -> nn.Module:
    return CellLocNet(meta)


def loss_fn(logits: torch.Tensor, offset_pred: torch.Tensor,
            soft_target: torch.Tensor, offset_target: torch.Tensor) -> torch.Tensor:
    """Soft cross-entropy against a bilinear tent distribution over cell
    centres + smooth-L1 on a residual offset target. soft_target is a dense
    [B, n_cells] distribution (already includes label smoothing);
    offset_target is [B, 2] in (-0.5, 0.5) cell units, the residual the
    pooled-centroid decode leaves. No confidence loss term, as before.
    """
    ce = -(soft_target * F.log_softmax(logits, dim=1)).sum(dim=1).mean()
    return ce + 1.0 * F.smooth_l1_loss(offset_pred, offset_target)


def export_onnx(model: nn.Module, path: str):
    model.eval()
    dummy = torch.zeros(1, 3, 128, 128)
    torch.onnx.export(model, dummy, path, input_names=["frame"],
                      output_names=["uvc"], opset_version=17, dynamo=False)


def estimate_position(frame: np.ndarray, onnx_path: str, meta: dict):
    """Deployment-shaped inference: frame (128x128x3 uint8) -> (lat, lon, conf).

    On-device this maps to the P4 runtime; here it runs the same ONNX artifact.
    """
    import onnxruntime as ort
    from pipeline.common import px_to_lonlat
    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    x = frame.astype(np.float32).transpose(2, 0, 1)[None] / 255.0
    u, v, conf = sess.run(None, {"frame": x})[0][0]
    lon, lat = px_to_lonlat(meta, float(u) * meta["width"], float(v) * meta["height"])
    return lat, lon, float(conf)

"""AGENT-EDITABLE — training procedure. One model per area (CLAUDE.md §1).

Current: map-cell classification with FULL-LATTICE coverage. Each EPOCH is one
complete pass over ALL train-lattice positions (not a random subsample), with a
fresh random rotation drawn per position per epoch, so over the run the model
sees epochs*len(crops) distinct (position, heading) vantages and every map cell
is visited many times — the coverage a C-way cell classifier needs to memorize
its box. Positions are streamed in shuffled chunks of CHUNK_SIZE so only one
chunk of pixels is ever resident. Targets are DECODE-CONSISTENT: the cell
target is a bilinear tent over the <=4 nearest cell centres (whose centroid is
the true position, matching what the 3x3 pooled decode reads), and the offset
target is the residual that pooled centroid leaves rather than the containing
cell's within-cell position. `--max-crops-per-bucket` is still accepted for
CLI compatibility but no longer controls sampling; it is logged as
`max_crops_per_bucket_arg`.

Usage: python -m model.train --area berlin --out-dir runs/<id> [--epochs 2]
"""

import argparse
import json
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

from model.model import CELL_PX, build_model, export_onnx, grid_dims, loss_fn
from pipeline.common import DATA_DIR, LIGHTING_BUCKETS, area_dir, load_meta
from pipeline.dataset import extract_crop, list_crops

CHUNK_SIZE = 4096  # lattice positions extracted at a time (bounds memory)
BATCH_SIZE = 64


def load_scene(area: str, data_dir: Path):
    """Meta + the relight image(s). One pass-through bucket on this branch."""
    meta = load_meta(area, data_dir)
    imgs = [np.asarray(Image.open(area_dir(area, data_dir) / "relight" / f"{b}.png"))
            for b in LIGHTING_BUCKETS]
    crops = list_crops(area, meta["width"], meta["height"], "train")
    return meta, imgs, crops


def extract_chunk(meta, img, crops, indices, rng):
    """Extract one chunk of vantages; return (x_uint8, cell_idx, pos).

    `cell_idx` is the containing cell, kept only for the top1 log line — the
    loss targets are built from `pos`, the continuous raster-pixel position.
    """
    gw, _ = grid_dims(meta)
    xs, cells, pos = [], [], []
    for i in indices:
        c = crops[i]
        cx, cy = c["cx"], c["cy"]
        angle = float(rng.uniform(0, 360))  # heading augmentation
        xs.append(extract_crop(img, cx, cy, angle))
        cells.append((cy // CELL_PX) * gw + (cx // CELL_PX))
        pos.append((float(cx), float(cy)))
    x = torch.from_numpy(np.stack(xs)).permute(0, 3, 1, 2).contiguous()
    return (x, torch.tensor(cells, dtype=torch.int64),
            torch.tensor(pos, dtype=torch.float32))


def make_soft_targets(pos, meta, n_cells, gw, gh, smoothing=0.05):
    """Bilinear tent target over the <=4 cell centres nearest each true
    position. pos: [B,2] raster px. Returns dense [B, n_cells] float32.

    The decode reads a 3x3-pooled centroid, so the target is the distribution
    whose centroid IS the true position: nearly identical border views now get
    nearly identical targets instead of contradictory hard labels. Index
    clamping merges out-of-grid neighbours' mass into the border cell via
    scatter_add_ — that mass is kept, not dropped, so rows always sum to 1.
    """
    fu = pos[:, 0] / CELL_PX - 0.5   # position in cell-centre coordinates
    fv = pos[:, 1] / CELL_PX - 0.5
    i0 = fu.floor(); j0 = fv.floor()
    a = (fu - i0); b = (fv - j0)     # in [0,1)
    tgt = torch.zeros(len(pos), n_cells, dtype=torch.float32, device=pos.device)
    for di, wu in ((0, 1 - a), (1, a)):
        for dj, wv in ((0, 1 - b), (1, b)):
            ci = (i0 + di).clamp(0, gw - 1).long()
            cj = (j0 + dj).clamp(0, gh - 1).long()
            tgt.scatter_add_(1, (cj * gw + ci).unsqueeze(1), (wu * wv).unsqueeze(1))
    return tgt * (1 - smoothing) + smoothing / n_cells


def residual_offset_target(model, logits, pos, meta):
    """(true position - model's 3x3 pooled centroid), in cell units,
    clamped to the offset head's (-0.5, 0.5) range. Detached.

    Mirrors forward()'s pooling exactly, using the model's own buffers so the
    grid geometry stays single-sourced. Early in training the argmax is noisy,
    so this target is nonstationary; that is expected and bounded by the clamp.
    """
    with torch.no_grad():
        p = F.softmax(logits, dim=1)
        idx = p.argmax(dim=1)
        win_u = model.cell_u.index_select(0, idx)
        win_v = model.cell_v.index_select(0, idx)
        du = (model.cell_u.unsqueeze(0) - win_u.unsqueeze(1)).abs()
        dv = (model.cell_v.unsqueeze(0) - win_v.unsqueeze(1)).abs()
        mask = ((du <= 1.6 * model.scale_u) & (dv <= 1.6 * model.scale_v)).float()
        q = p * mask
        m = q.sum(dim=1).clamp_min(1e-9)
        cu = (q * model.cell_u.unsqueeze(0)).sum(dim=1) / m
        cv = (q * model.cell_v.unsqueeze(0)).sum(dim=1) / m
        true_u = pos[:, 0] / meta["width"]; true_v = pos[:, 1] / meta["height"]
        ru = ((true_u - cu) / model.scale_u).clamp(-0.5, 0.5)
        rv = ((true_v - cv) / model.scale_v).clamp(-0.5, 0.5)
        return torch.stack([ru, rv], dim=1)


def train_area(area: str, out_dir: Path, data_dir: Path, epochs: int,
               max_crops_per_bucket: int, seed: int) -> dict:
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    device = "mps" if torch.backends.mps.is_available() else \
             "cuda" if torch.cuda.is_available() else "cpu"
    t0 = time.time()
    meta, imgs, crops = load_scene(area, data_dir)
    gw, gh = grid_dims(meta)
    model = build_model(meta).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    # A 12x-longer schedule wants a decaying LR, not the constant 1e-3.
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs,
                                                       eta_min=1e-4)
    print(f"[{area}] {len(crops)} lattice positions, {model.n_cells} cells, "
          f"device={device}")
    final_top1 = 0.0
    for epoch in range(epochs):
        model.train()
        loss_sum, correct, conf_sum, seen = 0.0, 0, 0.0, 0
        for img in imgs:
            # One full pass over the whole lattice, shuffled, fresh rotations.
            order = rng.permutation(len(crops))
            for start in range(0, len(order), CHUNK_SIZE):
                x_u8, cell_idx, pos = extract_chunk(
                    meta, img, crops, order[start:start + CHUNK_SIZE], rng)
                n = len(x_u8)
                perm = torch.randperm(n)
                for i in range(0, n, BATCH_SIZE):
                    idx = perm[i:i + BATCH_SIZE]
                    xb = x_u8[idx].to(device).float() / 255.0
                    cb, pb = cell_idx[idx].to(device), pos[idx].to(device)
                    logits, off_pred = model.training_outputs(xb)
                    soft_t = make_soft_targets(pb, meta, model.n_cells, gw, gh)
                    off_t = residual_offset_target(model, logits, pb, meta)
                    loss = loss_fn(logits, off_pred, soft_t, off_t)
                    opt.zero_grad(); loss.backward(); opt.step()
                    with torch.no_grad():
                        loss_sum += loss.item() * len(idx)
                        # top1 is vs. the CONTAINING cell, so it now reads lower
                        # than exp 4's 0.83 BY DESIGN: border views are supposed
                        # to split mass across the tent's cells. Judge
                        # convergence by loss and by eval, not by this.
                        correct += int((logits.argmax(dim=1) == cb).sum())
                        conf_sum += float(F.softmax(logits, dim=1).max(dim=1).values.sum())
                        seen += len(idx)
                del x_u8, cell_idx, pos
        sched.step()
        final_top1 = correct / max(seen, 1)
        print(f"[{area}] epoch {epoch + 1}/{epochs} loss={loss_sum / max(seen, 1):.4f} "
              f"top1={final_top1:.4f} mean_max_softmax={conf_sum / max(seen, 1):.4f} "
              f"views={seen} lr={sched.get_last_lr()[0]:.2e}")

    models_dir = out_dir / "models"
    models_dir.mkdir(parents=True, exist_ok=True)
    onnx_path = models_dir / f"{area}.onnx"
    export_onnx(model.cpu(), str(onnx_path))
    return {
        "area": area,
        "n_train_crops": epochs * len(crops) * len(imgs),
        "epochs": epochs,
        "device": device,
        "train_seconds": round(time.time() - t0, 1),
        "onnx_bytes": onnx_path.stat().st_size,
        "init": "from-scratch",  # §9: log init strategy per experiment
        "n_cells": model.n_cells,
        "distinct_vantages_seen": epochs * len(crops) * len(imgs),
        "full_lattice_epochs": True,
        "final_train_top1": round(final_top1, 4),
        "chunk_size": CHUNK_SIZE,
        "max_crops_per_bucket_arg": max_crops_per_bucket,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--area", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--data-dir", default=None)
    ap.add_argument("--epochs", type=int, default=2)
    ap.add_argument("--max-crops-per-bucket", type=int, default=800)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    data_dir = Path(args.data_dir) if args.data_dir else DATA_DIR
    out_dir = Path(args.out_dir)
    info = train_area(args.area, out_dir, data_dir, args.epochs,
                      args.max_crops_per_bucket, args.seed)
    out_dir.mkdir(parents=True, exist_ok=True)
    train_log = out_dir / "train_info.json"
    logs = json.loads(train_log.read_text()) if train_log.exists() else []
    logs.append(info)
    train_log.write_text(json.dumps(logs, indent=2))
    print(json.dumps(info))


if __name__ == "__main__":
    main()

"""AGENT-EDITABLE — model architecture + deployment inference wrapper.

The autoresearch loop may rewrite anything here (architecture, loss,
confidence mechanism, unified vs. dispatcher+specialists, pretrained init
per CLAUDE.md §3/§9) as long as:
  - train.py exports one ONNX per area at <out>/models/<area>.onnx taking a
    1x3x128x128 float32 input in [0,1] and returning [[u, v, conf]] with
    u, v normalized map coords and conf in [0,1];
  - the export passes pipeline/score.py's frozen deployment gates.

Current: localization as MAP-CELL CLASSIFICATION, from-scratch init, decoded by
LOCAL BELIEF POOLING. The map is cut into CELL_PX-square cells (Berlin:
55x54 = 2970); a conv trunk feeds a C-way softmax over cells plus a within-cell
offset head. The deployed decode pools belief over the 3x3 block of cells
centred on the argmax cell: confidence is the softmax mass on that whole
neighborhood, and the fix is the mass-weighted centroid of those cell centres,
nudged by the offset. Eval views stand 11-17 m off-lattice, so a view near a
cell border honestly splits its mass between adjacent cells — that is label
ambiguity, not place ambiguity, and pooling makes confidence measure "does the
model know WHERE it is" rather than "which arbitrary tile label applies".
Belief scattered across DISTANT cells still falls below threshold, so genuine
far-field ambiguity remains a safe abstention rather than a false fix. Cell
size is chosen so the parameterization alone suffices: a 128 m cell's
half-diagonal is 90.5 m, inside the 100 m usable radius, so naming the right
cell is already a usable fix.

The whole decode (softmax -> argmax -> neighborhood mask -> pooled centroid ->
offset -> conf) lives INSIDE the exported graph, because the scorer runs the
ONNX directly.
"""

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

CELL_PX = 128  # map cell edge in raster px (== metres at this GSD)


def grid_dims(meta: dict) -> tuple[int, int]:
    """(gw, gh): map-cell grid covering the area's raster."""
    return (math.ceil(meta["width"] / CELL_PX), math.ceil(meta["height"] / CELL_PX))


class CellLocNet(nn.Module):
    """Conv trunk -> GAP -> (C-way cell logits, within-cell offset).

    forward() is the DEPLOYED path and returns [B, 3] = (u, v, conf).
    training_outputs() exposes the raw (logits, offset) for the loss.
    """

    def __init__(self, meta: dict):
        super().__init__()
        gw, gh = grid_dims(meta)
        self.gw, self.gh = gw, gh
        self.n_cells = gw * gh
        width, height = meta["width"], meta["height"]

        chans = [3, 32, 64, 128, 160]
        layers = []
        for cin, cout in zip(chans, chans[1:]):
            layers += [nn.Conv2d(cin, cout, 3, stride=2, padding=1),
                       nn.BatchNorm2d(cout), nn.ReLU(inplace=True)]
        self.features = nn.Sequential(*layers)
        self.cls_head = nn.Linear(chans[-1], self.n_cells)
        self.off_head = nn.Linear(chans[-1], 2)

        # Cell-centre lookup tables, in normalized map coords. Registered as
        # buffers so they bake into the ONNX graph as constants.
        idx = np.arange(self.n_cells)
        cell_u = ((idx % gw) + 0.5) * CELL_PX / width
        cell_v = ((idx // gw) + 0.5) * CELL_PX / height
        self.register_buffer("cell_u", torch.tensor(cell_u, dtype=torch.float32))
        self.register_buffer("cell_v", torch.tensor(cell_v, dtype=torch.float32))
        self.scale_u = CELL_PX / width
        self.scale_v = CELL_PX / height

    def training_outputs(self, x):
        f = self.features(x).mean(dim=(2, 3))
        # tanh*0.5 keeps offsets inside the cell, in (-0.5, 0.5) cell units.
        return self.cls_head(f), torch.tanh(self.off_head(f)) * 0.5

    def forward(self, x):
        # Local belief pooling: confidence is the softmax mass on the 3x3 block
        # of cells centred on the argmax cell, and the fix is that block's
        # mass-weighted centroid (plus the within-cell offset). A border view
        # splitting mass between adjacent cells still names one place, so it
        # stays a fix; mass scattered across distant cells stays an abstention.
        logits, offset = self.training_outputs(x)
        p = F.softmax(logits, dim=1)
        idx = p.argmax(dim=1)
        win_u = self.cell_u.index_select(0, idx)
        win_v = self.cell_v.index_select(0, idx)
        # 1.6 cell-pitches per axis (Chebyshev) selects exactly the 3x3 block
        # with float-safety margin — the next ring sits at 2.0 pitches. Edge
        # cells simply have fewer neighbours in range; nothing is counted twice.
        du = (self.cell_u.unsqueeze(0) - win_u.unsqueeze(1)).abs()
        dv = (self.cell_v.unsqueeze(0) - win_v.unsqueeze(1)).abs()
        mask = ((du <= 1.6 * self.scale_u) & (dv <= 1.6 * self.scale_v)).float()
        q = p * mask
        conf = q.sum(dim=1)  # sub-sum of a softmax: in [0,1], >= single-cell max
        w = q / conf.unsqueeze(1).clamp_min(1e-9)
        u = (w * self.cell_u.unsqueeze(0)).sum(dim=1) + offset[:, 0] * self.scale_u
        v = (w * self.cell_v.unsqueeze(0)).sum(dim=1) + offset[:, 1] * self.scale_v
        return torch.stack([u.clamp(0.0, 1.0), v.clamp(0.0, 1.0), conf], dim=1)


def build_model(meta: dict) -> nn.Module:
    return CellLocNet(meta)


def loss_fn(logits: torch.Tensor, offset_pred: torch.Tensor,
            target_cell_idx: torch.Tensor, target_offset: torch.Tensor) -> torch.Tensor:
    """Cross-entropy on the true map cell + smooth-L1 on the within-cell offset.

    No confidence loss term: confidence IS the softmax mass on the winning
    cell, so honest abstention falls out of the parameterization rather than
    being trained against a proxy target. Label smoothing is light (0.05) so a
    memorized cell can still reach ~0.95 and clear the scorer's 0.3 threshold.
    """
    return (F.cross_entropy(logits, target_cell_idx, label_smoothing=0.05)
            + 1.0 * F.smooth_l1_loss(offset_pred, target_offset))


def export_onnx(model: nn.Module, path: str):
    model.eval()
    dummy = torch.zeros(1, 3, 128, 128)
    torch.onnx.export(model, dummy, path, input_names=["frame"],
                      output_names=["uvc"], opset_version=17, dynamo=False)


def estimate_position(frame: np.ndarray, onnx_path: str, meta: dict):
    """Deployment-shaped inference: frame (128x128x3 uint8) -> (lat, lon, conf).

    On-device this maps to the P4 runtime; here it runs the same ONNX artifact.
    """
    import onnxruntime as ort
    from pipeline.common import px_to_lonlat
    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    x = frame.astype(np.float32).transpose(2, 0, 1)[None] / 255.0
    u, v, conf = sess.run(None, {"frame": x})[0][0]
    lon, lat = px_to_lonlat(meta, float(u) * meta["width"], float(v) * meta["height"])
    return lat, lon, float(conf)

"""AGENT-EDITABLE — training procedure. One model per area (CLAUDE.md §1).

Current: map-cell classification with FULL-LATTICE coverage. Each EPOCH is one
complete pass over ALL train-lattice positions (not a random subsample), with a
fresh random rotation drawn per position per epoch, so over the run the model
sees epochs*len(crops) distinct (position, heading) vantages and every map cell
is visited many times — the coverage a C-way cell classifier needs to memorize
its box. Positions are streamed in shuffled chunks of CHUNK_SIZE so only one
chunk of pixels is ever resident. `--max-crops-per-bucket` is still accepted for
CLI compatibility but no longer controls sampling; it is logged as
`max_crops_per_bucket_arg`.

Usage: python -m model.train --area berlin --out-dir runs/<id> [--epochs 2]
"""

import argparse
import json
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

from model.model import CELL_PX, build_model, export_onnx, grid_dims, loss_fn
from pipeline.common import DATA_DIR, LIGHTING_BUCKETS, area_dir, load_meta
from pipeline.dataset import extract_crop, list_crops

CHUNK_SIZE = 4096  # lattice positions extracted at a time (bounds memory)
BATCH_SIZE = 64


def load_scene(area: str, data_dir: Path):
    """Meta + the relight image(s). One pass-through bucket on this branch."""
    meta = load_meta(area, data_dir)
    imgs = [np.asarray(Image.open(area_dir(area, data_dir) / "relight" / f"{b}.png"))
            for b in LIGHTING_BUCKETS]
    crops = list_crops(area, meta["width"], meta["height"], "train")
    return meta, imgs, crops


def extract_chunk(meta, img, crops, indices, rng):
    """Extract one chunk of vantages; return (x_uint8, cell_idx, offset)."""
    gw, _ = grid_dims(meta)
    xs, cells, offs = [], [], []
    for i in indices:
        c = crops[i]
        cx, cy = c["cx"], c["cy"]
        angle = float(rng.uniform(0, 360))  # heading augmentation
        xs.append(extract_crop(img, cx, cy, angle))
        cells.append((cy // CELL_PX) * gw + (cx // CELL_PX))
        offs.append(((cx % CELL_PX) / CELL_PX - 0.5,
                     (cy % CELL_PX) / CELL_PX - 0.5))
    x = torch.from_numpy(np.stack(xs)).permute(0, 3, 1, 2).contiguous()
    return (x, torch.tensor(cells, dtype=torch.int64),
            torch.tensor(offs, dtype=torch.float32))


def train_area(area: str, out_dir: Path, data_dir: Path, epochs: int,
               max_crops_per_bucket: int, seed: int) -> dict:
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    device = "mps" if torch.backends.mps.is_available() else \
             "cuda" if torch.cuda.is_available() else "cpu"
    t0 = time.time()
    meta, imgs, crops = load_scene(area, data_dir)
    model = build_model(meta).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    # A 12x-longer schedule wants a decaying LR, not the constant 1e-3.
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs,
                                                       eta_min=1e-4)
    print(f"[{area}] {len(crops)} lattice positions, {model.n_cells} cells, "
          f"device={device}")
    final_top1 = 0.0
    for epoch in range(epochs):
        model.train()
        loss_sum, correct, conf_sum, seen = 0.0, 0, 0.0, 0
        for img in imgs:
            # One full pass over the whole lattice, shuffled, fresh rotations.
            order = rng.permutation(len(crops))
            for start in range(0, len(order), CHUNK_SIZE):
                x_u8, cell_idx, off = extract_chunk(
                    meta, img, crops, order[start:start + CHUNK_SIZE], rng)
                n = len(x_u8)
                perm = torch.randperm(n)
                for i in range(0, n, BATCH_SIZE):
                    idx = perm[i:i + BATCH_SIZE]
                    xb = x_u8[idx].to(device).float() / 255.0
                    cb, ob = cell_idx[idx].to(device), off[idx].to(device)
                    logits, off_pred = model.training_outputs(xb)
                    loss = loss_fn(logits, off_pred, cb, ob)
                    opt.zero_grad(); loss.backward(); opt.step()
                    with torch.no_grad():
                        loss_sum += loss.item() * len(idx)
                        correct += int((logits.argmax(dim=1) == cb).sum())
                        conf_sum += float(F.softmax(logits, dim=1).max(dim=1).values.sum())
                        seen += len(idx)
                del x_u8, cell_idx, off
        sched.step()
        final_top1 = correct / max(seen, 1)
        print(f"[{area}] epoch {epoch + 1}/{epochs} loss={loss_sum / max(seen, 1):.4f} "
              f"top1={final_top1:.4f} mean_max_softmax={conf_sum / max(seen, 1):.4f} "
              f"views={seen} lr={sched.get_last_lr()[0]:.2e}")

    models_dir = out_dir / "models"
    models_dir.mkdir(parents=True, exist_ok=True)
    onnx_path = models_dir / f"{area}.onnx"
    export_onnx(model.cpu(), str(onnx_path))
    return {
        "area": area,
        "n_train_crops": epochs * len(crops) * len(imgs),
        "epochs": epochs,
        "device": device,
        "train_seconds": round(time.time() - t0, 1),
        "onnx_bytes": onnx_path.stat().st_size,
        "init": "from-scratch",  # §9: log init strategy per experiment
        "n_cells": model.n_cells,
        "distinct_vantages_seen": epochs * len(crops) * len(imgs),
        "full_lattice_epochs": True,
        "final_train_top1": round(final_top1, 4),
        "chunk_size": CHUNK_SIZE,
        "max_crops_per_bucket_arg": max_crops_per_bucket,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--area", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--data-dir", default=None)
    ap.add_argument("--epochs", type=int, default=2)
    ap.add_argument("--max-crops-per-bucket", type=int, default=800)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    data_dir = Path(args.data_dir) if args.data_dir else DATA_DIR
    out_dir = Path(args.out_dir)
    info = train_area(args.area, out_dir, data_dir, args.epochs,
                      args.max_crops_per_bucket, args.seed)
    out_dir.mkdir(parents=True, exist_ok=True)
    train_log = out_dir / "train_info.json"
    logs = json.loads(train_log.read_text()) if train_log.exists() else []
    logs.append(info)
    train_log.write_text(json.dumps(logs, indent=2))
    print(json.dumps(info))


if __name__ == "__main__":
    main()

"""AGENT-EDITABLE — model architecture + deployment inference wrapper.

The autoresearch loop may rewrite anything here (architecture, loss,
confidence mechanism, unified vs. dispatcher+specialists, pretrained init
per CLAUDE.md §3/§9) as long as:
  - train.py exports one ONNX per area at <out>/models/<area>.onnx taking a
    1x3x128x128 float32 input in [0,1] and returning [[u, v, conf]] with
    u, v normalized map coords and conf in [0,1];
  - the export passes pipeline/score.py's frozen deployment gates.

Current: localization as MAP-CELL CLASSIFICATION, from-scratch init. The map
is cut into CELL_PX-square cells (Berlin: 55x54 = 2970); a conv trunk feeds a
C-way softmax over cells plus a within-cell offset head. The deployed fix is
the winning cell's centre nudged by the offset, and confidence is the softmax
mass on that winning cell — so belief spread across many cells reads as low
confidence and the frame becomes a safe abstention rather than a false fix.
Cell size is chosen so the parameterization alone suffices: a 128 m cell's
half-diagonal is 90.5 m, inside the 100 m usable radius, so naming the right
cell is already a usable fix.

The whole decode (argmax -> cell-centre table lookup -> offset -> conf) lives
INSIDE the exported graph, because the scorer runs the ONNX directly.
"""

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

CELL_PX = 128  # map cell edge in raster px (== metres at this GSD)


def grid_dims(meta: dict) -> tuple[int, int]:
    """(gw, gh): map-cell grid covering the area's raster."""
    return (math.ceil(meta["width"] / CELL_PX), math.ceil(meta["height"] / CELL_PX))


class CellLocNet(nn.Module):
    """Conv trunk -> GAP -> (C-way cell logits, within-cell offset).

    forward() is the DEPLOYED path and returns [B, 3] = (u, v, conf).
    training_outputs() exposes the raw (logits, offset) for the loss.
    """

    def __init__(self, meta: dict):
        super().__init__()
        gw, gh = grid_dims(meta)
        self.gw, self.gh = gw, gh
        self.n_cells = gw * gh
        width, height = meta["width"], meta["height"]

        chans = [3, 32, 64, 128, 160]
        layers = []
        for cin, cout in zip(chans, chans[1:]):
            layers += [nn.Conv2d(cin, cout, 3, stride=2, padding=1),
                       nn.BatchNorm2d(cout), nn.ReLU(inplace=True)]
        self.features = nn.Sequential(*layers)
        self.cls_head = nn.Linear(chans[-1], self.n_cells)
        self.off_head = nn.Linear(chans[-1], 2)

        # Cell-centre lookup tables, in normalized map coords. Registered as
        # buffers so they bake into the ONNX graph as constants.
        idx = np.arange(self.n_cells)
        cell_u = ((idx % gw) + 0.5) * CELL_PX / width
        cell_v = ((idx // gw) + 0.5) * CELL_PX / height
        self.register_buffer("cell_u", torch.tensor(cell_u, dtype=torch.float32))
        self.register_buffer("cell_v", torch.tensor(cell_v, dtype=torch.float32))
        self.scale_u = CELL_PX / width
        self.scale_v = CELL_PX / height

    def training_outputs(self, x):
        f = self.features(x).mean(dim=(2, 3))
        # tanh*0.5 keeps offsets inside the cell, in (-0.5, 0.5) cell units.
        return self.cls_head(f), torch.tanh(self.off_head(f)) * 0.5

    def forward(self, x):
        logits, offset = self.training_outputs(x)
        p = F.softmax(logits, dim=1)
        conf, idx = p.max(dim=1)
        u = self.cell_u.index_select(0, idx) + offset[:, 0] * self.scale_u
        v = self.cell_v.index_select(0, idx) + offset[:, 1] * self.scale_v
        return torch.stack([u.clamp(0.0, 1.0), v.clamp(0.0, 1.0), conf], dim=1)


def build_model(meta: dict) -> nn.Module:
    return CellLocNet(meta)


def loss_fn(logits: torch.Tensor, offset_pred: torch.Tensor,
            target_cell_idx: torch.Tensor, target_offset: torch.Tensor) -> torch.Tensor:
    """Cross-entropy on the true map cell + smooth-L1 on the within-cell offset.

    No confidence loss term: confidence IS the softmax mass on the winning
    cell, so honest abstention falls out of the parameterization rather than
    being trained against a proxy target. Label smoothing is light (0.05) so a
    memorized cell can still reach ~0.95 and clear the scorer's 0.3 threshold.
    """
    return (F.cross_entropy(logits, target_cell_idx, label_smoothing=0.05)
            + 1.0 * F.smooth_l1_loss(offset_pred, target_offset))


def export_onnx(model: nn.Module, path: str):
    model.eval()
    dummy = torch.zeros(1, 3, 128, 128)
    torch.onnx.export(model, dummy, path, input_names=["frame"],
                      output_names=["uvc"], opset_version=17, dynamo=False)


def estimate_position(frame: np.ndarray, onnx_path: str, meta: dict):
    """Deployment-shaped inference: frame (128x128x3 uint8) -> (lat, lon, conf).

    On-device this maps to the P4 runtime; here it runs the same ONNX artifact.
    """
    import onnxruntime as ort
    from pipeline.common import px_to_lonlat
    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    x = frame.astype(np.float32).transpose(2, 0, 1)[None] / 255.0
    u, v, conf = sess.run(None, {"frame": x})[0][0]
    lon, lat = px_to_lonlat(meta, float(u) * meta["width"], float(v) * meta["height"])
    return lat, lon, float(conf)

"""AGENT-EDITABLE — training procedure. One model per area (CLAUDE.md §1).

Current: map-cell classification with FULL-LATTICE coverage. Each EPOCH is one
complete pass over ALL train-lattice positions (not a random subsample), with a
fresh random rotation drawn per position per epoch, so over the run the model
sees epochs*len(crops) distinct (position, heading) vantages and every map cell
is visited many times — the coverage a C-way cell classifier needs to memorize
its box. Positions are streamed in shuffled chunks of CHUNK_SIZE so only one
chunk of pixels is ever resident. `--max-crops-per-bucket` is still accepted for
CLI compatibility but no longer controls sampling; it is logged as
`max_crops_per_bucket_arg`.

Usage: python -m model.train --area berlin --out-dir runs/<id> [--epochs 2]
"""

import argparse
import json
import time
from pathlib import Path

import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image

from model.model import CELL_PX, build_model, export_onnx, grid_dims, loss_fn
from pipeline.common import DATA_DIR, LIGHTING_BUCKETS, area_dir, load_meta
from pipeline.dataset import extract_crop, list_crops

CHUNK_SIZE = 4096  # lattice positions extracted at a time (bounds memory)
BATCH_SIZE = 64


def load_scene(area: str, data_dir: Path):
    """Meta + the relight image(s). One pass-through bucket on this branch."""
    meta = load_meta(area, data_dir)
    imgs = [np.asarray(Image.open(area_dir(area, data_dir) / "relight" / f"{b}.png"))
            for b in LIGHTING_BUCKETS]
    crops = list_crops(area, meta["width"], meta["height"], "train")
    return meta, imgs, crops


def extract_chunk(meta, img, crops, indices, rng):
    """Extract one chunk of vantages; return (x_uint8, cell_idx, offset)."""
    gw, _ = grid_dims(meta)
    xs, cells, offs = [], [], []
    for i in indices:
        c = crops[i]
        cx, cy = c["cx"], c["cy"]
        angle = float(rng.uniform(0, 360))  # heading augmentation
        xs.append(extract_crop(img, cx, cy, angle))
        cells.append((cy // CELL_PX) * gw + (cx // CELL_PX))
        offs.append(((cx % CELL_PX) / CELL_PX - 0.5,
                     (cy % CELL_PX) / CELL_PX - 0.5))
    x = torch.from_numpy(np.stack(xs)).permute(0, 3, 1, 2).contiguous()
    return (x, torch.tensor(cells, dtype=torch.int64),
            torch.tensor(offs, dtype=torch.float32))


def train_area(area: str, out_dir: Path, data_dir: Path, epochs: int,
               max_crops_per_bucket: int, seed: int) -> dict:
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    device = "mps" if torch.backends.mps.is_available() else \
             "cuda" if torch.cuda.is_available() else "cpu"
    t0 = time.time()
    meta, imgs, crops = load_scene(area, data_dir)
    model = build_model(meta).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    # A 12x-longer schedule wants a decaying LR, not the constant 1e-3.
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs,
                                                       eta_min=1e-4)
    print(f"[{area}] {len(crops)} lattice positions, {model.n_cells} cells, "
          f"device={device}")
    final_top1 = 0.0
    for epoch in range(epochs):
        model.train()
        loss_sum, correct, conf_sum, seen = 0.0, 0, 0.0, 0
        for img in imgs:
            # One full pass over the whole lattice, shuffled, fresh rotations.
            order = rng.permutation(len(crops))
            for start in range(0, len(order), CHUNK_SIZE):
                x_u8, cell_idx, off = extract_chunk(
                    meta, img, crops, order[start:start + CHUNK_SIZE], rng)
                n = len(x_u8)
                perm = torch.randperm(n)
                for i in range(0, n, BATCH_SIZE):
                    idx = perm[i:i + BATCH_SIZE]
                    xb = x_u8[idx].to(device).float() / 255.0
                    cb, ob = cell_idx[idx].to(device), off[idx].to(device)
                    logits, off_pred = model.training_outputs(xb)
                    loss = loss_fn(logits, off_pred, cb, ob)
                    opt.zero_grad(); loss.backward(); opt.step()
                    with torch.no_grad():
                        loss_sum += loss.item() * len(idx)
                        correct += int((logits.argmax(dim=1) == cb).sum())
                        conf_sum += float(F.softmax(logits, dim=1).max(dim=1).values.sum())
                        seen += len(idx)
                del x_u8, cell_idx, off
        sched.step()
        final_top1 = correct / max(seen, 1)
        print(f"[{area}] epoch {epoch + 1}/{epochs} loss={loss_sum / max(seen, 1):.4f} "
              f"top1={final_top1:.4f} mean_max_softmax={conf_sum / max(seen, 1):.4f} "
              f"views={seen} lr={sched.get_last_lr()[0]:.2e}")

    models_dir = out_dir / "models"
    models_dir.mkdir(parents=True, exist_ok=True)
    onnx_path = models_dir / f"{area}.onnx"
    export_onnx(model.cpu(), str(onnx_path))
    return {
        "area": area,
        "n_train_crops": epochs * len(crops) * len(imgs),
        "epochs": epochs,
        "device": device,
        "train_seconds": round(time.time() - t0, 1),
        "onnx_bytes": onnx_path.stat().st_size,
        "init": "from-scratch",  # §9: log init strategy per experiment
        "n_cells": model.n_cells,
        "distinct_vantages_seen": epochs * len(crops) * len(imgs),
        "full_lattice_epochs": True,
        "final_train_top1": round(final_top1, 4),
        "chunk_size": CHUNK_SIZE,
        "max_crops_per_bucket_arg": max_crops_per_bucket,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--area", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--data-dir", default=None)
    ap.add_argument("--epochs", type=int, default=2)
    ap.add_argument("--max-crops-per-bucket", type=int, default=800)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    data_dir = Path(args.data_dir) if args.data_dir else DATA_DIR
    out_dir = Path(args.out_dir)
    info = train_area(args.area, out_dir, data_dir, args.epochs,
                      args.max_crops_per_bucket, args.seed)
    out_dir.mkdir(parents=True, exist_ok=True)
    train_log = out_dir / "train_info.json"
    logs = json.loads(train_log.read_text()) if train_log.exists() else []
    logs.append(info)
    train_log.write_text(json.dumps(logs, indent=2))
    print(json.dumps(info))


if __name__ == "__main__":
    main()

"""AGENT-EDITABLE — model architecture + deployment inference wrapper.

The autoresearch loop may rewrite anything here (architecture, loss,
confidence mechanism, unified vs. dispatcher+specialists, pretrained init
per CLAUDE.md §3/§9) as long as:
  - train.py exports one ONNX per area at <out>/models/<area>.onnx taking a
    1x3x128x128 float32 input in [0,1] and returning [[u, v, conf]] with
    u, v normalized map coords and conf in [0,1];
  - the export passes pipeline/score.py's frozen deployment gates.

Current: localization as MAP-CELL CLASSIFICATION, from-scratch init. The map
is cut into CELL_PX-square cells (Berlin: 55x54 = 2970); a conv trunk feeds a
C-way softmax over cells plus a within-cell offset head. The deployed fix is
the winning cell's centre nudged by the offset, and confidence is the softmax
mass on that winning cell — so belief spread across many cells reads as low
confidence and the frame becomes a safe abstention rather than a false fix.
Cell size is chosen so the parameterization alone suffices: a 128 m cell's
half-diagonal is 90.5 m, inside the 100 m usable radius, so naming the right
cell is already a usable fix.

The whole decode (argmax -> cell-centre table lookup -> offset -> conf) lives
INSIDE the exported graph, because the scorer runs the ONNX directly.
"""

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

CELL_PX = 128  # map cell edge in raster px (== metres at this GSD)


def grid_dims(meta: dict) -> tuple[int, int]:
    """(gw, gh): map-cell grid covering the area's raster."""
    return (math.ceil(meta["width"] / CELL_PX), math.ceil(meta["height"] / CELL_PX))


class CellLocNet(nn.Module):
    """Conv trunk -> GAP -> (C-way cell logits, within-cell offset).

    forward() is the DEPLOYED path and returns [B, 3] = (u, v, conf).
    training_outputs() exposes the raw (logits, offset) for the loss.
    """

    def __init__(self, meta: dict):
        super().__init__()
        gw, gh = grid_dims(meta)
        self.gw, self.gh = gw, gh
        self.n_cells = gw * gh
        width, height = meta["width"], meta["height"]

        chans = [3, 32, 64, 128, 160]
        layers = []
        for cin, cout in zip(chans, chans[1:]):
            layers += [nn.Conv2d(cin, cout, 3, stride=2, padding=1),
                       nn.BatchNorm2d(cout), nn.ReLU(inplace=True)]
        self.features = nn.Sequential(*layers)
        self.cls_head = nn.Linear(chans[-1], self.n_cells)
        self.off_head = nn.Linear(chans[-1], 2)

        # Cell-centre lookup tables, in normalized map coords. Registered as
        # buffers so they bake into the ONNX graph as constants.
        idx = np.arange(self.n_cells)
        cell_u = ((idx % gw) + 0.5) * CELL_PX / width
        cell_v = ((idx // gw) + 0.5) * CELL_PX / height
        self.register_buffer("cell_u", torch.tensor(cell_u, dtype=torch.float32))
        self.register_buffer("cell_v", torch.tensor(cell_v, dtype=torch.float32))
        self.scale_u = CELL_PX / width
        self.scale_v = CELL_PX / height

    def training_outputs(self, x):
        f = self.features(x).mean(dim=(2, 3))
        # tanh*0.5 keeps offsets inside the cell, in (-0.5, 0.5) cell units.
        return self.cls_head(f), torch.tanh(self.off_head(f)) * 0.5

    def forward(self, x):
        logits, offset = self.training_outputs(x)
        p = F.softmax(logits, dim=1)
        conf, idx = p.max(dim=1)
        u = self.cell_u.index_select(0, idx) + offset[:, 0] * self.scale_u
        v = self.cell_v.index_select(0, idx) + offset[:, 1] * self.scale_v
        return torch.stack([u.clamp(0.0, 1.0), v.clamp(0.0, 1.0), conf], dim=1)


def build_model(meta: dict) -> nn.Module:
    return CellLocNet(meta)


def loss_fn(logits: torch.Tensor, offset_pred: torch.Tensor,
            target_cell_idx: torch.Tensor, target_offset: torch.Tensor) -> torch.Tensor:
    """Cross-entropy on the true map cell + smooth-L1 on the within-cell offset.

    No confidence loss term: confidence IS the softmax mass on the winning
    cell, so honest abstention falls out of the parameterization rather than
    being trained against a proxy target. Label smoothing is light (0.05) so a
    memorized cell can still reach ~0.95 and clear the scorer's 0.3 threshold.
    """
    return (F.cross_entropy(logits, target_cell_idx, label_smoothing=0.05)
            + 1.0 * F.smooth_l1_loss(offset_pred, target_offset))


def export_onnx(model: nn.Module, path: str):
    model.eval()
    dummy = torch.zeros(1, 3, 128, 128)
    torch.onnx.export(model, dummy, path, input_names=["frame"],
                      output_names=["uvc"], opset_version=17, dynamo=False)


def estimate_position(frame: np.ndarray, onnx_path: str, meta: dict):
    """Deployment-shaped inference: frame (128x128x3 uint8) -> (lat, lon, conf).

    On-device this maps to the P4 runtime; here it runs the same ONNX artifact.
    """
    import onnxruntime as ort
    from pipeline.common import px_to_lonlat
    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    x = frame.astype(np.float32).transpose(2, 0, 1)[None] / 255.0
    u, v, conf = sess.run(None, {"frame": x})[0][0]
    lon, lat = px_to_lonlat(meta, float(u) * meta["width"], float(v) * meta["height"])
    return lat, lon, float(conf)

"""AGENT-EDITABLE — training procedure. One model per area (CLAUDE.md §1).

Current: map-cell classification. Each EPOCH redraws a fresh random subset of
the full train lattice with fresh random rotations, so over the run the model
sees epochs*max_crops_per_bucket distinct (position, heading) vantages — the
coverage the C-way cell classifier needs, since one fixed draw of 6,000 crops
cannot cover ~3,000 cells.

Usage: python -m model.train --area berlin --out-dir runs/<id> [--epochs 2]
"""

import argparse
import json
import time
from pathlib import Path

import numpy as np
import torch
from PIL import Image

from model.model import CELL_PX, build_model, export_onnx, grid_dims, loss_fn
from pipeline.common import DATA_DIR, LIGHTING_BUCKETS, area_dir, load_meta
from pipeline.dataset import extract_crop, list_crops


def load_scene(area: str, data_dir: Path):
    """Meta + the relight image(s). One pass-through bucket on this branch."""
    meta = load_meta(area, data_dir)
    imgs = [np.asarray(Image.open(area_dir(area, data_dir) / "relight" / f"{b}.png"))
            for b in LIGHTING_BUCKETS]
    crops = list_crops(area, meta["width"], meta["height"], "train")
    return meta, imgs, crops


def sample_epoch(meta, imgs, crops, max_crops_per_bucket: int, rng):
    """Draw a fresh crop subset with fresh rotations; return (x, cell_idx, off)."""
    gw, _ = grid_dims(meta)
    xs, cells, offs = [], [], []
    for img in imgs:
        picks = rng.choice(len(crops), size=min(max_crops_per_bucket, len(crops)),
                           replace=False)
        for i in picks:
            c = crops[i]
            cx, cy = c["cx"], c["cy"]
            angle = float(rng.uniform(0, 360))  # heading augmentation
            xs.append(extract_crop(img, cx, cy, angle))
            cells.append((cy // CELL_PX) * gw + (cx // CELL_PX))
            offs.append(((cx % CELL_PX) / CELL_PX - 0.5,
                         (cy % CELL_PX) / CELL_PX - 0.5))
    x = torch.from_numpy(np.stack(xs)).permute(0, 3, 1, 2).contiguous().float() / 255.0
    return (x, torch.tensor(cells, dtype=torch.int64),
            torch.tensor(offs, dtype=torch.float32))


def train_area(area: str, out_dir: Path, data_dir: Path, epochs: int,
               max_crops_per_bucket: int, seed: int) -> dict:
    torch.manual_seed(seed)
    rng = np.random.default_rng(seed)
    device = "mps" if torch.backends.mps.is_available() else \
             "cuda" if torch.cuda.is_available() else "cpu"
    t0 = time.time()
    meta, imgs, crops = load_scene(area, data_dir)
    model = build_model(meta).to(device)
    opt = torch.optim.Adam(model.parameters(), lr=1e-3)
    n = 0
    print(f"[{area}] {len(crops)} lattice positions, {model.n_cells} cells, "
          f"device={device}")
    for epoch in range(epochs):
        # Fresh vantages every epoch, not one extract-once set reused.
        x, cell_idx, off = sample_epoch(meta, imgs, crops, max_crops_per_bucket, rng)
        n = len(x)
        model.train()
        perm = torch.randperm(n)
        losses = []
        for i in range(0, n, 64):
            idx = perm[i:i + 64]
            xb = x[idx].to(device)
            cb, ob = cell_idx[idx].to(device), off[idx].to(device)
            logits, off_pred = model.training_outputs(xb)
            loss = loss_fn(logits, off_pred, cb, ob)
            opt.zero_grad(); loss.backward(); opt.step()
            losses.append(loss.item())
        print(f"[{area}] epoch {epoch + 1}/{epochs} loss={np.mean(losses):.4f}")

    models_dir = out_dir / "models"
    models_dir.mkdir(parents=True, exist_ok=True)
    onnx_path = models_dir / f"{area}.onnx"
    export_onnx(model.cpu(), str(onnx_path))
    return {
        "area": area,
        "n_train_crops": n,
        "epochs": epochs,
        "device": device,
        "train_seconds": round(time.time() - t0, 1),
        "onnx_bytes": onnx_path.stat().st_size,
        "init": "from-scratch",  # §9: log init strategy per experiment
        "n_cells": model.n_cells,
        "distinct_vantages_seen": epochs * max_crops_per_bucket,
    }


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--area", required=True)
    ap.add_argument("--out-dir", required=True)
    ap.add_argument("--data-dir", default=None)
    ap.add_argument("--epochs", type=int, default=2)
    ap.add_argument("--max-crops-per-bucket", type=int, default=800)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()
    data_dir = Path(args.data_dir) if args.data_dir else DATA_DIR
    out_dir = Path(args.out_dir)
    info = train_area(args.area, out_dir, data_dir, args.epochs,
                      args.max_crops_per_bucket, args.seed)
    out_dir.mkdir(parents=True, exist_ok=True)
    train_log = out_dir / "train_info.json"
    logs = json.loads(train_log.read_text()) if train_log.exists() else []
    logs.append(info)
    train_log.write_text(json.dumps(logs, indent=2))
    print(json.dumps(info))


if __name__ == "__main__":
    main()