#!/usr/bin/env python3
"""One-time: extract prebackbone weights from a full YOLO best.pt checkpoint."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import torch

# Allow running before vendor is installed
_HERE = Path(__file__).resolve().parent
for base in (_HERE / "vendor", _HERE.parent / "ultralytics"):
    if (base / "ultralytics" / "__init__.py").exists():
        sys.path.insert(0, str(base))
        break


def _clean_state_dict(state: dict) -> dict:
    """Drop thop profiling keys (total_ops, total_params) not in nn.Module."""
    skip = ("total_ops", "total_params")
    return {k: v for k, v in state.items() if not any(k == s or k.endswith(f".{s}") for s in skip)}


def extract(
    full_ckpt: Path,
    out_path: Path,
    prebackbone_name: str | None = None,
) -> Path:
    from ultralytics.nn.tasks import torch_safe_load

    ckpt, _ = torch_safe_load(str(full_ckpt))
    train_args = ckpt.get("train_args") or {}
    name = (prebackbone_name or train_args.get("prebackbone") or "A11_CA").upper()
    channels = int(train_args.get("channels", 3) or 3)

    model = (ckpt.get("ema") or ckpt.get("model")).float()
    pb = getattr(model, "prebackbone", None)
    if pb is None:
        raise RuntimeError(f"No prebackbone in {full_ckpt}")

    state = _clean_state_dict(pb.state_dict())

    # Save only tensors + metadata (no YOLO backbone/head)
    payload = {
        "prebackbone": name,
        "channels": channels,
        "state_dict": state,
        "source_checkpoint": str(full_ckpt.resolve()),
    }
    out_path.parent.mkdir(parents=True, exist_ok=True)
    torch.save(payload, out_path)

    n_params = sum(t.numel() for t in payload["state_dict"].values())
    size_mb = out_path.stat().st_size / (1024 * 1024)
    print(f"Saved {out_path} ({size_mb:.2f} MB, {n_params:,} parameters, type={name})")
    return out_path


def main() -> None:
    p = argparse.ArgumentParser(description="Extract prebackbone-only weights from best.pt")
    p.add_argument(
        "--ckpt",
        type=Path,
        default=_HERE / "weights" / "best.pt",
        help="Full YOLO checkpoint (best.pt)",
    )
    p.add_argument(
        "--out",
        type=Path,
        default=_HERE / "weights" / "prebackbone_a11_ca.pt",
        help="Output path for prebackbone-only weights",
    )
    p.add_argument("--name", type=str, default=None, help="Override prebackbone type (default: from train_args)")
    args = p.parse_args()
    if not args.ckpt.exists():
        raise SystemExit(f"Checkpoint not found: {args.ckpt}")
    extract(args.ckpt, args.out, args.name)


if __name__ == "__main__":
    main()