"""Fine-tune YOLO11 on the unified dataset.

YOLO11n is the primary arm and YOLO11s the second, which is the "lightweight
backbone comparison" for free. Both are trained by one invocation so the runs
sit side by side under the same project directory.

The training recipe itself is stock Ultralytics, identical to the reference
YOLO11 fine-tuning notebook: of the 105 resolved config keys, only `batch`,
`device`, `project`, `name` and `exist_ok` differ, and all 55 recipe keys --
optimiser, learning-rate schedule, loss gains, and every augmentation -- match.
This script is plumbing around `YOLO.train()`, not a modified training loop.

Reports per-class mAP, not just the mean. The mean hides the case this project
actually cares about: cones are the smallest class and stop signs come from a
different domain, so either can be quietly terrible while the average looks fine.

Note: Ultralytics YOLO11 weights are AGPL-3.0.

Usage:
    python -m src.train.train --models yolo11n.pt yolo11s.pt --epochs 100
    python -m src.train.train --wandb          # + train/val loss and grad norm
"""

from __future__ import annotations

import argparse
import os
from pathlib import Path

import pandas as pd

from src.common import paths, schema


def find_manifest(root: Path, data_yaml: Path) -> Path | None:
    """The manifest sits beside the data.yaml when pulled from the Hub, and
    under the unified root when built locally. Try both."""
    for candidate in (data_yaml.parent / "manifest.parquet", Path(root) / "manifest.parquet"):
        if candidate.exists():
            return candidate
    return None


def print_class_balance(manifest_path: Path | None) -> None:
    """Instance counts before training, because the imbalance is severe.

    Barriers are everywhere in nuScenes, stop signs come almost entirely from
    COCO. Knowing the counts up front is what makes a
    weak per-class mAP readable afterwards.
    """
    if manifest_path is None:
        print("(no manifest found, skipping the class balance report)")
        return

    # This is a diagnostic, not a prerequisite. A dataset pulled from the Hub is
    # perfectly trainable without a readable manifest, so a failure here must not
    # take the training run down with it.
    try:
        objects = schema.objects_only(schema.read_manifest(manifest_path))
        print("\ntraining instances per class")
        print(pd.crosstab(objects["split"], objects["class"]).to_string())
        print()
    except Exception as error:                # noqa: BLE001 -- report and carry on
        print(f"(could not read {manifest_path}: {error}; skipping the balance report)")


def print_per_class_metrics(metrics, names: dict[int, str]) -> None:
    print("\n" + "-" * 62)
    print(f"{'class':<14}{'images':>8}{'instances':>11}{'mAP50':>9}{'mAP50-95':>11}")
    print("-" * 62)
    for index, class_id in enumerate(metrics.box.ap_class_index):
        precision, recall, ap50, ap = metrics.box.class_result(index)
        print(f"{names.get(int(class_id), class_id):<14}{'':>8}{'':>11}"
              f"{ap50:>9.3f}{ap:>11.3f}")
    print("-" * 62)
    print(f"{'all':<14}{'':>8}{'':>11}{metrics.box.map50:>9.3f}{metrics.box.map:>11.3f}")
    print("-" * 62)

    missing = [names[i] for i in names if i not in set(int(c) for c in metrics.box.ap_class_index)]
    if missing:
        print(f"!! no val instances for: {', '.join(missing)}")


# Ultralytics normalises the optimiser step to a nominal batch of 64 by
# accumulating gradients: accumulate = round(nbs / batch), and weight decay is
# rescaled by batch * accumulate / nbs to compensate. That compensation is exact
# only when the batch divides 64. At batch 37 -- the kind of number AutoBatch
# (batch=-1) picks from free GPU memory -- effective weight decay lands 15.6%
# above the reference recipe, and at 48 it lands 25% below.
NOMINAL_BATCH = 64


def weight_decay_scale(batch: int) -> float:
    """How far this batch size moves effective weight decay from the default."""
    accumulate = max(round(NOMINAL_BATCH / batch), 1)
    return batch * accumulate / NOMINAL_BATCH


def warn_if_batch_perturbs_recipe(batch: int) -> None:
    if batch < 0:
        print("!! --batch -1 lets AutoBatch pick from free GPU memory. It returns an")
        print("!! arbitrary integer, so effective weight decay drifts from the")
        print("!! reference recipe by as much as 25-56%. Prefer 16, 32 or 64.")
        return
    scale = weight_decay_scale(batch)
    if abs(scale - 1.0) > 0.01:
        print(f"!! batch={batch} does not divide {NOMINAL_BATCH}, so effective weight")
        print(f"!! decay is {scale:.3f}x the configured value ({(scale - 1) * 100:+.1f}%).")
        print(f"!! Use a divisor of {NOMINAL_BATCH} to match the stock recipe exactly.")


def train_one(model_name: str, data_yaml: Path, args) -> None:
    from ultralytics import YOLO

    print("=" * 62)
    print(f"training {model_name}")
    print("=" * 62)

    model = YOLO(model_name)

    close_logger = None
    if args.wandb:
        from src.train import wandb_logger

        close_logger = wandb_logger.attach(
            model,
            project=args.wandb_project,
            run_name=f"{Path(model_name).stem}-{args.epochs}ep",
            extra_config={"dataset": str(data_yaml)},
            log_grad_norm=args.log_grad_norm,
        )
        print(f"logging to W&B project {args.wandb_project!r}"
              + (" (with gradient norm)" if args.log_grad_norm else ""))

    try:
        _run_training(model, model_name, data_yaml, args)
    finally:
        # Ultralytics fires on_train_end only on success, and the gradient-norm
        # probe patches a torch function that must be put back either way --
        # otherwise a crash training the first model leaves a wrapper in place
        # for the second.
        if close_logger is not None:
            close_logger()

    metrics = model.val(data=str(data_yaml), imgsz=args.imgsz, device=args.device)
    print_per_class_metrics(metrics, model.names)
    print(f"\nweights: {args.project}/{Path(model_name).stem}/weights/best.pt")


def _run_training(model, model_name: str, data_yaml: Path, args) -> None:
    model.train(
        data=str(data_yaml),
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=args.batch,
        device=args.device,
        workers=args.workers,
        project=str(args.project),
        name=f"{Path(model_name).stem}",
        exist_ok=True,
        seed=args.seed,
    )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--unified-root", type=Path, default=None)
    parser.add_argument("--data", type=Path, default=None,
                        help="defaults to <unified-root>/yolo/data.yaml")
    parser.add_argument("--models", nargs="+", default=["yolo11n.pt", "yolo11s.pt"])
    parser.add_argument("--epochs", type=int, default=100)
    parser.add_argument("--imgsz", type=int, default=640)
    # Ultralytics' own default, and what the reference notebook trains with.
    # Divides NOMINAL_BATCH, so the optimiser recipe matches exactly.
    parser.add_argument("--batch", type=int, default=16)
    parser.add_argument("--device", default="0")
    parser.add_argument("--workers", type=int, default=8)
    parser.add_argument("--project", type=Path, default=Path("runs/detect"))
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--wandb", action="store_true",
                        help="log train/val loss and gradient norm to Weights & Biases; "
                             "authenticate with $WANDB_API_KEY")
    parser.add_argument("--wandb-project", default="cone-distance")
    parser.add_argument("--log-grad-norm", action="store_true",
                        help="also log gradient norm. Needs --wandb. Unlike the loss "
                             "callbacks, this patches torch.nn.utils.clip_grad_norm_ "
                             "for the duration of training, because Ultralytics "
                             "discards the value and fires no callback while the "
                             "gradients are still live.")
    args = parser.parse_args()

    if args.wandb and not os.environ.get("WANDB_API_KEY"):
        print("!! --wandb given but $WANDB_API_KEY is not set; wandb will prompt "
              "or fall back to offline mode")

    root = paths.unified_root(args.unified_root)
    data_yaml = args.data or (root / "yolo" / "data.yaml")
    if not data_yaml.exists():
        raise SystemExit(f"{data_yaml} does not exist; run src.data.to_yolo first")

    warn_if_batch_perturbs_recipe(args.batch)
    print_class_balance(find_manifest(root, data_yaml))
    for model_name in args.models:
        train_one(model_name, data_yaml, args)


if __name__ == "__main__":
    main()