"""Prepare, train, evaluate, and predict. See README.md for commands."""

import argparse
import json
import os
import random
import time
from pathlib import Path

# Keep downloaded weights with the project, not in a global cache.
os.environ.setdefault("HF_HOME", str(Path(".cache/huggingface").resolve()))

import numpy as np
import torch
from torch import nn
from torch.utils.data import ConcatDataset, DataLoader, RandomSampler, WeightedRandomSampler

from data import (
    BACKGROUND,
    IMAGE_NAMES,
    FolderDataset,
    GridDataset,
    SupervisedGrids,
    grid_boxes,
    prepare,
    prepare_jsonl,
    prepare_folders,
    read_rgb,
    record_image_path,
    target_name,
    transform,
)
from models import BACKBONE, GridModel, select_target, transfer_backbone


def save_json(path, value):
    Path(path).write_text(json.dumps(value, indent=2) + "\n")


def device_for(name):
    if name == "auto":
        name = (
            "cuda"
            if torch.cuda.is_available()
            else "mps"
            if torch.backends.mps.is_available()
            else "cpu"
        )
    if name == "mps" and not torch.backends.mps.is_available():
        raise RuntimeError(
            "MPS is unavailable. Run in a terminal with Apple GPU access; use --device cpu only intentionally."
        )
    if name == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA is unavailable on this machine.")
    return torch.device(name)


def empty_cache(device):
    if device.type == "mps":
        torch.mps.empty_cache()
    elif device.type == "cuda":
        torch.cuda.empty_cache()


def seed_everything(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.set_num_threads(4)


def records_for(manifest, tiles, split):
    return [r for r in manifest["records"] if r["tiles"] == tiles and r["split"] == split]


def auxiliary_records(manifest, targets, split):
    return [r for r in records_for(manifest, 16, split) if r["target"] in targets]


def folder_records(folders, kind, split):
    return (
        [r for r in folders["records"] if r["kind"] == kind and r["split"] == split]
        if folders
        else []
    )


def cell_logits(model, images):
    """Logits shaped (batch, classes, cells); a separate tile is one cell."""
    output = model(images)
    return output.unsqueeze(-1) if model.tiles == 9 else output.flatten(2)


def partial_loss(logits, labels, weights, any_labels, any_weights):
    """Weighted BCE on known cells, plus presence anywhere (max over cells) for image-level labels."""
    cells = nn.functional.binary_cross_entropy_with_logits(
        logits, labels, weight=weights, reduction="sum"
    )
    anywhere = nn.functional.binary_cross_entropy_with_logits(
        logits.amax(-1), any_labels, weight=any_weights, reduction="sum"
    )
    total = weights.sum() + any_weights.sum()
    return (cells + anywhere) / total.clamp(min=1), total


@torch.inference_mode()
def evaluate_folders(model, loader, device):
    model.eval()
    loss, total, probabilities, labels, weights = 0.0, 0.0, [], [], []
    for images, *supervision in loader:
        supervision = [x.to(device) for x in supervision]
        logits = cell_logits(model, images.to(device)).float()
        batch_loss, batch_total = partial_loss(logits, *supervision)
        loss += batch_loss.item() * batch_total.item()
        total += batch_total.item()
        probabilities.append(logits.amax(-1).sigmoid().cpu())
        labels.append(supervision[2].cpu())
        weights.append(supervision[3].cpu())
    return (
        loss / max(total, 1),
        torch.cat(probabilities),
        torch.cat(labels).bool(),
        torch.cat(weights),
    )


def folder_report(probabilities, labels, weights, classes, threshold):
    """Single-label folder images: recall per class, and how often each class fires on other classes' images."""
    predicted = probabilities >= threshold
    background = (weights == 1).all(1) & ~labels.any(1)
    report = {
        "threshold": threshold,
        "images": len(labels),
        "per_class": {},
        "background_images": int(background.sum()),
        "background_any_selected": float(predicted[background].any(1).float().mean())
        if background.any()
        else None,
    }
    for i, name in enumerate(classes):
        positive, others = labels[:, i], ~labels[:, i] & labels.any(1)
        report["per_class"][name] = {
            "positives": int(positive.sum()),
            "recall": float(predicted[positive, i].float().mean()) if positive.any() else None,
            "fires_on_other_classes": float(predicted[others, i].float().mean())
            if others.any()
            else None,
        }
    return report


def thresholds_for(checkpoint, targets):
    return torch.tensor(
        [
            checkpoint.get("thresholds", {}).get(
                checkpoint["classes"][int(i)], checkpoint["threshold"]
            )
            for i in targets
        ]
    ).unsqueeze(1)


def grid_logits(model, images, targets):
    if model.tiles == 9:
        batch, cells = images.shape[:2]
        logits = model(images.flatten(0, 1)).reshape(batch, cells, -1).transpose(1, 2)
    else:
        logits = model(images).flatten(2)
    selected = select_target(logits, targets)
    # Presence in any enlarged subregion counts as presence in the tile.
    for target_id in getattr(model, "crop_target_ids", []):
        mask = targets == target_id
        if not mask.any():
            continue
        tile_images = images[mask].flatten(0, 1)
        height, width = tile_images.shape[-2:]
        crop_h, crop_w = height // 2, width // 2
        positions = [
            (0, 0),
            (0, width - crop_w),
            (height - crop_h, 0),
            (height - crop_h, width - crop_w),
            ((height - crop_h) // 2, (width - crop_w) // 2),
        ]
        scores = selected[mask]
        for top, left in positions:
            crop = tile_images[:, :, top : top + crop_h, left : left + crop_w]
            crop = nn.functional.interpolate(
                crop, size=(height, width), mode="bilinear", align_corners=False
            )
            enlarged = model(crop)[:, target_id].reshape_as(scores)
            scores = torch.maximum(scores, enlarged)
        selected[mask] = scores
    return selected


@torch.inference_mode()
def evaluate_loader(model, loader, device):
    model.eval()
    predictions, labels, targets = [], [], []
    loss, total = 0.0, 0
    for images, target, truth in loader:
        logits = grid_logits(model, images.to(device), target.to(device))
        loss += nn.functional.binary_cross_entropy_with_logits(
            logits, truth.to(device), reduction="sum"
        ).item()
        total += truth.numel()
        predictions.append(logits.sigmoid().cpu())
        labels.append(truth)
        targets.append(target)
    if not total:
        raise ValueError("Evaluation split is empty")
    return loss / total, torch.cat(predictions), torch.cat(labels).bool(), torch.cat(targets)


def metrics(probabilities, labels, threshold):
    predicted = probabilities >= threshold
    tp = int((predicted & labels).sum())
    fp = int((predicted & ~labels).sum())
    fn = int((~predicted & labels).sum())
    return {
        "grids": len(labels),
        "exact_grid_accuracy": float((predicted == labels).all(1).float().mean()),
        "cell_accuracy": float((predicted == labels).float().mean()),
        "cell_precision": tp / max(tp + fp, 1),
        "cell_recall": tp / max(tp + fn, 1),
        "cell_f1": 2 * tp / max(2 * tp + fp + fn, 1),
        "tp": tp,
        "fp": fp,
        "fn": fn,
    }


def choose_threshold(probabilities, labels):
    candidates = [round(0.1 + 0.025 * i, 3) for i in range(33)]

    def rank(threshold):
        result = metrics(probabilities, labels, threshold)
        return result["exact_grid_accuracy"], result["cell_f1"], -abs(threshold - 0.5)

    return max(candidates, key=rank)


def load_checkpoint(path, device):
    checkpoint = torch.load(path, map_location="cpu", weights_only=True)
    model = GridModel(
        checkpoint["tiles"],
        len(checkpoint["classes"]),
        pretrained=False,
        backbone=checkpoint.get("backbone", BACKBONE),
    )
    model.load_state_dict(checkpoint["model"])
    model.crop_target_ids = [
        checkpoint["classes"].index(name) for name in checkpoint.get("crop_targets", [])
    ]
    if model.crop_target_ids and checkpoint["tiles"] != 9:
        raise ValueError("Crop views are supported only by tile classifiers")
    return model.to(device).eval(), checkpoint


def train_one(args, tiles, manifest, folders=None):
    seed_everything(args.seed)
    device = device_for(args.device)
    output = Path(args.output) / f"grid{tiles}"
    if (output / "best.pt").exists():
        raise ValueError(
            f"{output}/best.pt already exists. Choose a new --output to preserve this run."
        )
    output.mkdir(parents=True, exist_ok=True)
    cells = 1 if tiles == 9 else 16
    kind = "tile" if tiles == 9 else "full"
    classes = sorted(
        {r["target"] for r in manifest["records"] if r["tiles"] == tiles}
        | {t for r in folder_records(folders, kind, "train") for t in r["targets"]} - {BACKGROUND}
    )
    aux_targets = sorted({target_name(name) for name in args.aux_targets})
    if aux_targets and (tiles != 9 or not set(aux_targets) <= set(classes)):
        raise ValueError("Auxiliary targets must be existing classes in a 9-tile model")
    # 4x4 photos are 450 px; 448 keeps small slivers at a size divisible by the stride.
    size = args.size or (224 if tiles == 9 else 448)
    backbone_lr = args.backbone_lr if args.backbone_lr is not None else 0.0001
    finetune_head_lr = args.finetune_head_lr if args.finetune_head_lr is not None else 0.0003
    config = {
        "tiles": tiles,
        "classes": classes,
        "size": size,
        "backbone": args.backbone,
        "seed": args.seed,
        "batch_size": args.batch_size,
        "device": str(device),
        "precision": "bfloat16 autocast" if args.amp else "float32",
        "frozen_epochs": args.frozen_epochs,
        "finetune_epochs": args.epochs,
        "patience": args.patience,
        "optimizer": "AdamW",
        "weight_decay": 0.0001,
        "head_lr": args.head_lr,
        "finetune_head_lr": finetune_head_lr,
        "finetune_backbone_lr": backbone_lr,
        "warm_start": args.warm_start,
        "init_backbone": args.init_backbone,
        "aux_targets": aux_targets,
        "folders": args.folders,
        "folder_kind": kind if folders else None,
        "folder_negative_weight": args.negative_weight if folders else None,
        "old_fraction": args.old_fraction if folders else None,
        "epoch_size": args.epoch_size,
        "keep_uncertain": args.keep_uncertain,
        "train_batchnorm": not args.freeze_bn,
        "amp": args.amp,
        "selection": (
            "mean of answer-key and folder validation loss"
            if folders
            else "mean of main and auxiliary validation loss"
            if aux_targets
            else "main validation loss"
        ),
        "torch": str(torch.__version__),
    }
    if args.amp and device.type != "cuda":
        raise ValueError("--amp is supported only with --device cuda")
    if device.type == "cuda":
        torch.backends.cudnn.benchmark = True
        torch.set_float32_matmul_precision("high")
    save_json(output / "config.json", config)
    save_json(output / "classes.json", classes)
    save_json(output / "split.json", manifest)
    training_records = records_for(manifest, tiles, "train") + auxiliary_records(
        manifest, aux_targets, "train"
    )
    training = SupervisedGrids(
        GridDataset(
            manifest["root"],
            training_records,
            classes,
            size,
            training=True,
            tiles=tiles == 9,
            include_full_images=bool(aux_targets),
        ),
        cells,
        keep_uncertain=args.keep_uncertain,
        flip=not args.no_flip,
    )
    validation = GridDataset(manifest["root"], records_for(manifest, tiles, "val"), classes, size)
    workers = {
        "num_workers": args.workers,
        "persistent_workers": args.workers > 0,
        "pin_memory": device.type == "cuda",
    }
    sampler, folder_loader = None, None
    if folders:
        folder_training = FolderDataset(
            folders["root"],
            folder_records(folders, kind, "train"),
            classes,
            size,
            cells,
            training=True,
            negative_weight=args.negative_weight,
        )
        folder_validation = FolderDataset(
            folders["root"],
            folder_records(folders, kind, "val"),
            classes,
            size,
            cells,
            negative_weight=args.negative_weight,
        )
        # Square-root class balancing: rare classes are seen more often without memorizing a few images.
        counts = {}
        for record in folder_training.records:
            counts[record["targets"][0]] = counts.get(record["targets"][0], 0) + 1
        folder_weights = torch.tensor(
            [counts[r["targets"][0]] ** -0.5 for r in folder_training.records]
        )
        folder_weights *= (1 - args.old_fraction) / folder_weights.sum()
        old_weights = torch.full((len(training),), args.old_fraction / len(training))
        sampler = WeightedRandomSampler(
            torch.cat([old_weights, folder_weights]).double(),
            args.epoch_size or len(training) + len(folder_training),
            generator=torch.Generator().manual_seed(args.seed),
        )
        training = ConcatDataset([training, folder_training])
        folder_loader = DataLoader(folder_validation, batch_size=64, **workers)
    elif args.epoch_size:
        sampler = RandomSampler(
            training,
            replacement=True,
            num_samples=args.epoch_size,
            generator=torch.Generator().manual_seed(args.seed),
        )
    train_loader = DataLoader(
        training, batch_size=args.batch_size, shuffle=sampler is None, sampler=sampler, **workers
    )
    val_loader = DataLoader(validation, batch_size=16, **workers)
    aux_loader = None
    if aux_targets:
        aux_validation = GridDataset(
            manifest["root"],
            auxiliary_records(manifest, aux_targets, "val"),
            classes,
            size,
            tile_grids=True,
        )
        if not len(aux_validation):
            raise ValueError("Auxiliary targets need validation examples")
        aux_loader = DataLoader(aux_validation, batch_size=2)
    if args.warm_start:
        model, initial = load_checkpoint(args.warm_start, device)
        config["backbone"] = initial.get("backbone", BACKBONE)
        save_json(output / "config.json", config)
        if initial["tiles"] != tiles or initial["classes"] != classes or initial["size"] != size:
            raise ValueError("Warm-start model must match task, classes, and image size")
    else:
        model = GridModel(tiles, len(classes), backbone=args.backbone).to(device)
    if args.init_backbone:
        source = torch.load(args.init_backbone, map_location="cpu", weights_only=True)
        print(
            f"grid{tiles}: copied {transfer_backbone(model, source['model'])} backbone tensors "
            f"from {args.init_backbone}",
            flush=True,
        )
    model.freeze_backbone(True)
    optimizer = torch.optim.AdamW(model.head.parameters(), lr=args.head_lr, weight_decay=0.0001)
    history, best, stale = [], float("inf"), 0
    start = time.perf_counter()
    print(
        f"grid{tiles}: {len(train_loader.dataset)} training items ({len(train_loader.sampler)} per epoch), "
        f"{len(validation)} validation grids, device={device}",
        flush=True,
    )
    for epoch in range(args.frozen_epochs + args.epochs):
        frozen = epoch < args.frozen_epochs
        if epoch == args.frozen_epochs:
            model.freeze_backbone(False)
            optimizer = torch.optim.AdamW(
                [
                    {"params": model.backbone.parameters(), "lr": backbone_lr},
                    {"params": model.head.parameters(), "lr": finetune_head_lr},
                ],
                weight_decay=0.0001,
            )
            stale = 0
        model.train()
        if args.freeze_bn:
            # Keep pretrained batch-normalization statistics (for small datasets).
            model.backbone.eval()
        total_loss, supervised = 0.0, 0.0
        epoch_start = time.perf_counter()
        for images, *supervision in train_loader:
            images = images.to(device, non_blocking=True)
            supervision = [x.to(device, non_blocking=True) for x in supervision]
            optimizer.zero_grad(set_to_none=True)
            with torch.autocast("cuda", dtype=torch.bfloat16, enabled=args.amp):
                logits = cell_logits(model, images)
            loss, count = partial_loss(logits.float(), *supervision)
            loss.backward()
            optimizer.step()
            total_loss += loss.item() * count.item()
            supervised += count.item()
        val_loss, probabilities, truth, _ = evaluate_loader(model, val_loader, device)
        aux_loss = evaluate_loader(model, aux_loader, device)[0] if aux_loader is not None else None
        folder_loss = (
            evaluate_folders(model, folder_loader, device)[0] if folder_loader is not None else None
        )
        extra = [x for x in (aux_loss, folder_loss) if x is not None]
        selection_loss = sum([val_loss] + extra) / (1 + len(extra))
        row = {
            "epoch": epoch + 1,
            "phase": "head" if frozen else "finetune",
            "train_loss": total_loss / max(supervised, 1),
            "val_loss": val_loss,
            "aux_val_loss": aux_loss,
            "folder_val_loss": folder_loss,
            "selection_loss": selection_loss,
            "val_exact_at_0.5": metrics(probabilities, truth, 0.5)["exact_grid_accuracy"],
            "seconds": time.perf_counter() - epoch_start,
        }
        history.append(row)
        save_json(output / "history.json", history)
        print(json.dumps({"grid": tiles, **row}), flush=True)
        if selection_loss < best:
            best, stale = selection_loss, 0
            torch.save(
                {
                    **config,
                    "epoch": epoch + 1,
                    "val_loss": val_loss,
                    "selection_loss": selection_loss,
                    "model": {k: v.detach().cpu().clone() for k, v in model.state_dict().items()},
                },
                output / "best.pt",
            )
        elif not frozen:
            stale += 1
        if not frozen and stale >= args.patience:
            break
    elapsed = time.perf_counter() - start
    del optimizer, model
    empty_cache(device)
    model, checkpoint = load_checkpoint(output / "best.pt", device)
    _, probabilities, truth, val_targets = evaluate_loader(model, val_loader, device)
    threshold = choose_threshold(probabilities, truth)
    checkpoint["threshold"] = threshold
    checkpoint["thresholds"] = {}
    if aux_loader is not None:
        _, aux_probabilities, aux_truth, aux_ids = evaluate_loader(model, aux_loader, device)
        for name in aux_targets:
            mask = aux_ids == classes.index(name)
            if mask.sum() >= 3:
                checkpoint["thresholds"][name] = choose_threshold(
                    aux_probabilities[mask], aux_truth[mask]
                )
        save_json(
            output / "aux_validation.json",
            {
                "source": "4x4 crops; proxy only for 3x3 performance",
                "thresholds": checkpoint["thresholds"],
                **metrics(aux_probabilities, aux_truth, thresholds_for(checkpoint, aux_ids)),
            },
        )
    if folder_loader is not None:
        _, folder_probabilities, folder_truth, folder_weights = evaluate_folders(
            model, folder_loader, device
        )
        save_json(
            output / "folder_validation.json",
            {
                "source": f"held-out {kind} images from {args.folders}; single-label, so 'fires_on_other_classes' "
                "includes images where that target is genuinely visible",
                **folder_report(
                    folder_probabilities, folder_truth, folder_weights, classes, threshold
                ),
            },
        )
    checkpoint["training_seconds"] = elapsed
    torch.save(checkpoint, output / "best.pt")
    save_json(
        output / "validation.json",
        {
            "threshold": threshold,
            "thresholds": checkpoint["thresholds"],
            **metrics(probabilities, truth, thresholds_for(checkpoint, val_targets)),
        },
    )
    print(
        f"grid{tiles}: saved best epoch {checkpoint['epoch']}; threshold={threshold}; training={elapsed:.1f}s",
        flush=True,
    )


def synchronize(device):
    if device.type == "mps":
        torch.mps.synchronize()
    elif device.type == "cuda":
        torch.cuda.synchronize()


def calibrate_crops(args):
    """Enable enlarged views only after a validation-only comparison."""
    device = device_for(args.device)
    seed_everything(42)
    model, checkpoint = load_checkpoint(args.checkpoint, device)
    name = target_name(args.target)
    if checkpoint["tiles"] != 9 or name not in checkpoint["classes"]:
        raise ValueError("Choose a supported target in a tile classifier")
    if name in checkpoint.get("crop_targets", []):
        raise ValueError("Crop views are already enabled for this target")
    directory = Path(args.checkpoint).parent
    manifest = json.loads((directory / "split.json").read_text())
    records = auxiliary_records(manifest, [name], "val")
    if len(records) < 3:
        raise ValueError("Need at least three auxiliary validation grids")
    dataset = GridDataset(
        manifest["root"], records, checkpoint["classes"], checkpoint["size"], tile_grids=True
    )
    loader = DataLoader(dataset, batch_size=2)
    _, probabilities, truth, _ = evaluate_loader(model, loader, device)
    old_threshold = checkpoint.get("thresholds", {}).get(name, checkpoint["threshold"])
    before = metrics(probabilities, truth, old_threshold)
    model.crop_target_ids.append(checkpoint["classes"].index(name))
    _, probabilities, truth, _ = evaluate_loader(model, loader, device)
    cutoff = choose_threshold(probabilities, truth)
    after = metrics(probabilities, truth, cutoff)
    accepted = (
        after["exact_grid_accuracy"] >= before["exact_grid_accuracy"]
        and after["cell_f1"] > before["cell_f1"]
    )
    report = {
        "target": name,
        "source": "4x4 validation crops; not 3x3 test data",
        "views": "full tile plus center and four corner half-size crops",
        "before": {"threshold": old_threshold, **before},
        "after": {"threshold": cutoff, **after},
        "accepted": accepted,
    }
    if accepted:
        backup = directory / "pre_crop_best.pt"
        if backup.exists():
            raise ValueError(f"Refusing to overwrite {backup}")
        torch.save(checkpoint, backup)
        checkpoint.setdefault("crop_targets", []).append(name)
        checkpoint.setdefault("thresholds", {})[name] = cutoff
        torch.save(checkpoint, args.checkpoint)
    save_json(directory / "crop_calibration.json", report)
    print(json.dumps(report, indent=2))


@torch.inference_mode()
def benchmark(model, dataset, device):
    images, target, _ = dataset[0]
    images, target = images.unsqueeze(0).to(device), torch.tensor([target], device=device)
    for _ in range(5):
        grid_logits(model, images, target)
    synchronize(device)
    timings = []
    for _ in range(30):
        start = time.perf_counter()
        grid_logits(model, images, target).sigmoid().cpu()
        synchronize(device)
        timings.append((time.perf_counter() - start) * 1000)
    return {
        "median_ms": float(np.median(timings)),
        "p95_ms": float(np.percentile(timings, 95)),
        "scope": "one grid, warm model, includes output transfer; excludes image loading and preprocessing",
        "device": str(device),
        "iterations": len(timings),
    }


def evaluate(args):
    device = device_for(args.device)
    seed_everything(42)
    model, checkpoint = load_checkpoint(args.checkpoint, device)
    directory = Path(args.checkpoint).parent
    manifest = json.loads((directory / "split.json").read_text())
    records = records_for(manifest, checkpoint["tiles"], args.split)
    if not records:
        raise ValueError(
            f"No {args.split} grids in this manifest"
            + (" (run prepare with --review)" if args.split == "review" else "")
        )
    dataset = GridDataset(
        args.data or manifest["root"], records, checkpoint["classes"], checkpoint["size"]
    )
    loss, probabilities, truth, targets = evaluate_loader(
        model, DataLoader(dataset, batch_size=4), device
    )
    threshold = thresholds_for(checkpoint, targets)
    report = {
        "split": args.split,
        "epoch": checkpoint["epoch"],
        "threshold": checkpoint["threshold"],
        "thresholds": checkpoint.get("thresholds", {}),
        "loss": loss,
        "training_seconds": checkpoint["training_seconds"],
        "overall": metrics(probabilities, truth, threshold),
        "per_class": {},
    }
    if args.split == "review":
        # Same grids, answered by the labeler whose output the model was trained on.
        right = [r["labeler_answers"] == r["answers"] for r in records if "labeler_answers" in r]
        report["labeler_exact_grid_accuracy"] = sum(right) / len(right) if right else None
    for i, name in enumerate(checkpoint["classes"]):
        mask = targets == i
        support = {
            split: len(
                [
                    r
                    for r in records_for(manifest, checkpoint["tiles"], split)
                    if r["target"] == name
                ]
            )
            for split in ("train", "val", "test", "review")
        }
        report["per_class"][name] = {
            "support": support,
            "metrics": metrics(probabilities[mask], truth[mask], threshold[mask])
            if mask.any()
            else None,
        }
    report["latency"] = benchmark(model, dataset, device)
    save_json(directory / f"{args.split}_metrics.json", report)
    save_json(
        directory / f"{args.split}_predictions.json",
        [
            {
                "id": r["id"],
                "target": r["target"],
                "answers": r["answers"],
                "selected_cells": (p >= cutoff).nonzero().flatten().tolist(),
                "probabilities": p.tolist(),
            }
            for r, p, cutoff in zip(records, probabilities, threshold)
        ],
    )
    print(json.dumps(report, indent=2))


@torch.inference_mode()
def predict(args):
    device = device_for(args.device)
    model, checkpoint = load_checkpoint(args.checkpoint, device)
    name = target_name(args.target)
    if name not in checkpoint["classes"]:
        raise ValueError(f"Unknown target {name!r}; supported: {checkpoint['classes']}")
    preprocess = transform(checkpoint["size"])
    image = read_rgb(image_file(args.image))
    if checkpoint["tiles"] == 9:
        images = torch.stack([preprocess(image.crop(box)) for box in grid_boxes(*image.size, 9)])
    else:
        images = preprocess(image)
    target = torch.tensor([checkpoint["classes"].index(name)], device=device)
    probabilities = grid_logits(model, images.unsqueeze(0).to(device), target).sigmoid().cpu()[0]
    cutoff = checkpoint.get("thresholds", {}).get(name, checkpoint["threshold"])
    print(
        json.dumps(
            {
                "target": name,
                "threshold": cutoff,
                "selected_cells": (probabilities >= cutoff).nonzero().flatten().tolist(),
                "probabilities": probabilities.tolist(),
            },
            indent=2,
        )
    )


def export(args):
    """ONNX fp32 plus calibrated int8, checked against PyTorch on the old test grids."""
    import onnx
    from onnxruntime.quantization import (
        CalibrationDataReader,
        CalibrationMethod,
        QuantFormat,
        QuantType,
        quantize_static,
    )
    from onnxruntime.quantization.shape_inference import quant_pre_process

    from predict_onnx import GridPredictor

    seed_everything(42)
    model, checkpoint = load_checkpoint(args.checkpoint, torch.device("cpu"))
    if model.crop_target_ids:
        raise ValueError("Enlarged crop views are not supported by the ONNX export")
    tiles, size = checkpoint["tiles"], checkpoint["size"]
    directory = Path(args.checkpoint).parent
    output = Path(args.output or directory / "onnx")
    output.mkdir(parents=True, exist_ok=True)
    manifest = json.loads((directory / "split.json").read_text())
    root = Path(args.data or manifest["root"])

    class Probabilities(nn.Module):
        def __init__(self):
            super().__init__()
            self.model = model

        def forward(self, images):
            logits = self.model(images)
            return (logits if tiles == 9 else logits.flatten(2)).sigmoid()

    info = {
        "tiles": tiles,
        "size": size,
        "classes": checkpoint["classes"],
        "threshold": checkpoint["threshold"],
        "thresholds": checkpoint.get("thresholds", {}),
        "epoch": checkpoint["epoch"],
    }

    def with_metadata(path):
        exported = onnx.load(path)
        del exported.metadata_props[:]
        exported.metadata_props.add(key="grid_vision", value=json.dumps(info))
        onnx.save(exported, path)

    fp32, prepared, int8 = (
        output / f"grid{tiles}.onnx",
        output / f"grid{tiles}.pre.onnx",
        output / f"grid{tiles}.int8.onnx",
    )
    for path in (fp32, int8):
        if path.exists():
            raise ValueError(f"{path} already exists. Choose a new --output.")
    torch.onnx.export(
        Probabilities().eval(),
        torch.zeros(9 if tiles == 9 else 1, 3, size, size),
        fp32,
        input_names=["images"],
        output_names=["probabilities"],
        opset_version=17,
        dynamic_axes={"images": {0: "batch"}, "probabilities": {0: "batch"}},
        dynamo=False,
    )
    with_metadata(fp32)
    reference = GridPredictor(fp32)
    calibration = random.Random(0).sample(
        records_for(manifest, tiles, "train"),
        min(args.calibration_grids, len(records_for(manifest, tiles, "train"))),
    )

    class Calibration(CalibrationDataReader):
        def __init__(self):
            self.batches = iter(reference.inputs(record_image_path(root, r)) for r in calibration)

        def get_next(self):
            batch = next(self.batches, None)
            return None if batch is None else {"images": batch}

    quant_pre_process(str(fp32), str(prepared))
    # Percentile ranges and full-precision first/last layers changed the fewest baseline decisions.
    layers = [
        n.name for n in onnx.load(prepared).graph.node if n.op_type in ("Conv", "Gemm", "MatMul")
    ]
    quantize_static(
        str(prepared),
        str(int8),
        Calibration(),
        quant_format=QuantFormat.QDQ,
        per_channel=True,
        activation_type=QuantType.QUInt8,
        weight_type=QuantType.QInt8,
        calibrate_method=CalibrationMethod.Percentile,
        nodes_to_exclude=[layers[0], layers[-1]],
    )
    prepared.unlink()
    with_metadata(int8)

    records = records_for(manifest, tiles, "test")
    variants = {"fp32": GridPredictor(fp32), "int8": GridPredictor(int8)}
    probabilities = {name: [] for name in ["torch", *variants]}
    labels, cutoffs = [], []
    with torch.inference_mode():
        for record in records:
            inputs = reference.inputs(record_image_path(root, record))
            target = checkpoint["classes"].index(record["target"])
            scores = Probabilities()(torch.from_numpy(inputs)).numpy()
            probabilities["torch"].append(scores.T[target] if tiles == 9 else scores[0, target])
            for name, predictor in variants.items():
                probabilities[name].append(predictor.probabilities(inputs)[target])
            labels.append([i in record["answers"] for i in range(record["tiles"])])
            cutoffs.append(info["thresholds"].get(record["target"], info["threshold"]))
    labels, cutoffs = torch.tensor(labels), torch.tensor(cutoffs).unsqueeze(1)
    tensors = {name: torch.tensor(np.array(values)) for name, values in probabilities.items()}
    selected = {name: values >= cutoffs for name, values in tensors.items()}
    report = {
        "checkpoint": str(args.checkpoint),
        "test_grids": len(records),
        "preprocessing": "predict_onnx.preprocess (Pillow + numpy) for every variant",
    }
    for name, path in [("fp32", fp32), ("int8", int8)]:
        timings = {}
        sample = reference.inputs(record_image_path(root, records[0]))
        for threads in (1, 4):
            predictor = GridPredictor(path, threads)
            for _ in range(5):
                predictor.probabilities(sample)
            runs = []
            for _ in range(50):
                start = time.perf_counter()
                predictor.probabilities(sample)
                runs.append((time.perf_counter() - start) * 1000)
            timings[f"{threads}_thread_median_ms"] = float(np.median(runs))
        report[name] = {
            "file": str(path),
            "megabytes": path.stat().st_size / 1e6,
            "max_abs_difference_from_torch": float((tensors[name] - tensors["torch"]).abs().max()),
            "grids_with_different_selection": int(
                (selected[name] != selected["torch"]).any(1).sum()
            ),
            "test": metrics(tensors[name], labels, cutoffs),
            "cpu_latency_per_grid": {
                **timings,
                "scope": "one grid on this machine's CPU, excludes image loading",
            },
        }
    report["torch_test"] = metrics(tensors["torch"], labels, cutoffs)
    save_json(output / "export.json", report)
    print(json.dumps(report, indent=2))


def seam_ratio(image, parts=3):
    """Separate tiles have sharp seams at the 1/parts lines; one continuous photo usually does not."""
    pixels = np.asarray(image.convert("L"), dtype=np.float32)
    height, width = pixels.shape
    dx, dy = np.abs(np.diff(pixels, axis=1)).mean(0), np.abs(np.diff(pixels, axis=0)).mean(1)

    def seams(differences, length):
        return max(
            differences[max(0, round(length * j / parts) - 2) : round(length * j / parts) + 1].max()
            for j in range(1, parts)
        )

    return min(seams(dx, width), seams(dy, height)) / ((np.median(dx) + np.median(dy)) / 2 + 1e-6)


def image_file(source):
    """A challenge image, or the image inside a dataset item folder."""
    source = Path(source)
    if source.is_dir():
        found = next((source / name for name in IMAGE_NAMES if (source / name).exists()), None)
        if found is None:
            raise ValueError(f"No image in {source}")
        return found
    return source


def layout(source):
    """9 or 16 tiles for a challenge image or item folder (300 px = 3x3, 450 px = 4x4)."""
    metadata = Path(source) / "metadata.json"
    if Path(source).is_dir() and metadata.is_file():
        grid = json.loads(metadata.read_text()).get("grid")
        if grid in ("3x3", "4x4"):
            return 9 if grid == "3x3" else 16
    image = read_rgb(image_file(source))
    side = max(image.size)
    if abs(side - 300) <= 30:
        return 9
    if abs(side - 450) <= 45:
        return 16
    return 9 if seam_ratio(image) >= 3 else 16


def solve(args):
    """Detect the layout and use the matching grid model."""
    args.checkpoint = args.grid9 if layout(args.image) == 9 else args.grid16
    predict(args)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    prep = commands.add_parser("prepare")
    prep.add_argument(
        "--data",
        required=True,
        help="Folder of labeled <item>/{image.jpg, metadata.json} (or answer-key) folders",
    )
    prep.add_argument("--output", default="runs")
    prep.add_argument("--seed", type=int, default=42)
    prep.add_argument(
        "--workers", type=int, default=os.cpu_count() or 1, help="Processes for scanning images"
    )
    prep.add_argument(
        "--sam",
        action="store_true",
        help="Take 4x4 answers from <item>/sam.json (SAM 3) instead of metadata.json",
    )
    prep.add_argument(
        "--review",
        nargs="*",
        default=[],
        help="Reference answer files held out as a 'review' split",
    )
    prep_folders = commands.add_parser("prepare-folders", help="Ingest class-folder image datasets")
    prep_folders.add_argument("--data", nargs="+", required=True, help="Class-folder datasets")
    prep_folders.add_argument(
        "--manifest", default="runs/split.json", help="Old split; its held-out images are excluded"
    )
    prep_folders.add_argument("--output", default="runs")
    prep_folders.add_argument("--seed", type=int, default=42)
    train = commands.add_parser("train")
    train.add_argument("--manifest", default="runs/split.json")
    train.add_argument("--output", default="runs/baseline")
    train.add_argument("--tiles", choices=["9", "16", "both"], default="both")
    train.add_argument("--device", choices=["auto", "cuda", "mps", "cpu"], default="auto")
    train.add_argument("--batch-size", type=int, default=16)
    train.add_argument("--frozen-epochs", type=int, default=3)
    train.add_argument(
        "--epochs", type=int, default=20, help="Maximum fine-tuning epochs after head training"
    )
    train.add_argument("--patience", type=int, default=5)
    train.add_argument("--backbone-lr", type=float, help="Default: 1e-4")
    train.add_argument("--finetune-head-lr", type=float, help="Default: 3e-4")
    train.add_argument("--head-lr", type=float, default=0.001)
    train.add_argument("--warm-start", help="Initialize from a compatible trained checkpoint")
    train.add_argument(
        "--aux-targets",
        nargs="*",
        default=[],
        help="Use these targets' 4x4 crops to train a 9-tile classifier",
    )
    train.add_argument("--seed", type=int, default=42)
    train.add_argument("--size", type=int, help="Input size. Default: 224 for 9 tiles, 448 for 16")
    train.add_argument(
        "--folders", help="Folder manifest from prepare-folders, e.g. runs/folders.json"
    )
    train.add_argument(
        "--negative-weight",
        type=float,
        default=0.2,
        help="Weight for assuming other targets are absent in a single-label folder image",
    )
    train.add_argument(
        "--old-fraction", type=float, default=0.25, help="Share of samples from answer-key grids"
    )
    train.add_argument("--epoch-size", type=int, help="Training samples (tiles or grids) per epoch")
    train.add_argument(
        "--keep-uncertain",
        action="store_true",
        help="Also train on cells the labeler marked uncertain (ignored by default)",
    )
    train.add_argument(
        "--no-flip", action="store_true", help="Turn off random left-right mirroring"
    )
    train.add_argument(
        "--backbone",
        default=BACKBONE,
        help="Any timm model with a stride-16 stage, e.g. convnext_tiny.dinov3_lvd1689m",
    )
    train.add_argument(
        "--init-backbone", help="Copy backbone weights from another checkpoint, e.g. the 3x3 model"
    )
    train.add_argument(
        "--freeze-bn", action="store_true", help="Keep pretrained batch-normalization statistics"
    )
    train.add_argument(
        "--workers",
        type=int,
        default=min(16, os.cpu_count() or 1),
        help="DataLoader worker processes",
    )
    train.add_argument("--amp", action="store_true", help="bfloat16 autocast (CUDA only)")
    test = commands.add_parser("evaluate")
    test.add_argument("--checkpoint", required=True)
    test.add_argument(
        "--split",
        choices=["test", "review"],
        default="test",
        help="review: the reference grids held out by prepare --review",
    )
    test.add_argument("--data", help="Override dataset root if moved")
    test.add_argument("--device", choices=["auto", "cuda", "mps", "cpu"], default="auto")
    prediction = commands.add_parser("predict")
    prediction.add_argument("--checkpoint", required=True)
    prediction.add_argument("--image", required=True, help="Image file or example folder")
    prediction.add_argument("--target", required=True)
    prediction.add_argument("--device", choices=["auto", "cuda", "mps", "cpu"], default="auto")
    onnx_export = commands.add_parser(
        "export", help="ONNX fp32 and int8 for CPU inference without PyTorch"
    )
    onnx_export.add_argument("--checkpoint", required=True)
    onnx_export.add_argument("--output", help="Default: <checkpoint folder>/onnx")
    onnx_export.add_argument("--data", help="Override dataset root if moved")
    onnx_export.add_argument("--calibration-grids", type=int, default=100)
    solver = commands.add_parser("solve", help="Detect the layout and use the matching model")
    solver.add_argument("--image", required=True, help="Image file or example folder")
    solver.add_argument("--target", required=True)
    solver.add_argument("--grid9", default="runs/baseline/grid9/best.pt")
    solver.add_argument("--grid16", default="runs/baseline/grid16/best.pt")
    solver.add_argument("--device", choices=["auto", "cuda", "mps", "cpu"], default="cpu")
    crops = commands.add_parser("calibrate-crops")
    crops.add_argument("--checkpoint", required=True)
    crops.add_argument("--target", required=True)
    crops.add_argument("--device", choices=["auto", "cuda", "mps", "cpu"], default="auto")
    args = parser.parse_args()
    if args.command == "prepare":
        source = Path(args.data)
        if (source / "data" / "3x3").is_dir() or (source / "3x3").is_dir():
            if args.sam or args.review:
                parser.error("--sam and --review apply only to raw labeled folders")
            _, counts = prepare_jsonl(source, args.output)
        else:
            _, counts = prepare(args.data, args.output, args.seed, args.workers, args.sam, args.review)
        print(json.dumps(counts, indent=2))
    elif args.command == "prepare-folders":
        _, summary = prepare_folders(args.data, args.manifest, args.output, args.seed)
        print(json.dumps(summary, indent=2))
    elif args.command == "train":
        if min(args.batch_size, args.epochs, args.patience) < 1 or args.frozen_epochs < 0:
            parser.error(
                "Batch size, epochs, and patience must be positive; frozen epochs cannot be negative"
            )
        if any(
            lr is not None and (not np.isfinite(lr) or lr <= 0)
            for lr in (args.backbone_lr, args.finetune_head_lr, args.head_lr)
        ):
            parser.error("Learning rates must be finite and positive")
        if not 0 < args.old_fraction < 1 or args.negative_weight < 0 or args.workers < 0:
            parser.error("--old-fraction must be in (0, 1); weights and workers cannot be negative")
        if not np.isfinite(args.negative_weight):
            parser.error("--negative-weight must be finite")
        if (args.epoch_size is not None and args.epoch_size < 1) or (
            args.size is not None and args.size < 1
        ):
            parser.error("--epoch-size and --size must be positive")
        manifest = json.loads(Path(args.manifest).read_text())
        folders = json.loads(Path(args.folders).read_text()) if args.folders else None
        for tiles in [9, 16] if args.tiles == "both" else [int(args.tiles)]:
            train_one(args, tiles, manifest, folders)
    elif args.command == "evaluate":
        evaluate(args)
    elif args.command == "export":
        if args.calibration_grids < 1:
            parser.error("--calibration-grids must be positive")
        export(args)
    elif args.command == "solve":
        solve(args)
    elif args.command == "calibrate-crops":
        calibrate_crops(args)
    else:
        predict(args)


if __name__ == "__main__":
    main()