"""Model and hardware profile for the Methods section.

Measures parameter count, GFLOPs, checkpoint size, single-slice latency and
batched throughput for each backbone on the machine it is run on, and writes a
JSON/Markdown pair the manuscript can quote directly. Latency is timed over a
configurable number of consecutive axial slices of a real scan held in memory,
with CUDA synchronisation around every call, so the figure is reproducible and
the sample size is recorded alongside it rather than asserted in prose.

Usage
-----
    python scripts/06_profile.py [--n 200] [--weights path/to/best.pt]
"""

from __future__ import annotations

import argparse
import json
import platform
import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from luna_rev import config as cfg
from luna_rev.io_luna import all_uids, normalize_hu, read_volume


def measure_latency(model, images, device, warmup: int = 20) -> dict:
    """Per-slice latency at batch size 1, including pre/post-processing."""
    import torch

    for img in images[:warmup]:
        model.predict([img], imgsz=cfg.IMG_SIZE, conf=cfg.INFER_CONF,
                      max_det=cfg.INFER_MAX_DET, device=device, verbose=False)
    if torch.cuda.is_available():
        torch.cuda.synchronize()

    times = []
    for img in images:
        t0 = time.perf_counter()
        model.predict([img], imgsz=cfg.IMG_SIZE, conf=cfg.INFER_CONF,
                      max_det=cfg.INFER_MAX_DET, device=device, verbose=False)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        times.append((time.perf_counter() - t0) * 1000.0)
    a = np.array(times)
    return {"n": len(a), "mean_ms": float(a.mean()), "sd_ms": float(a.std(ddof=1)),
            "median_ms": float(np.median(a)), "p95_ms": float(np.percentile(a, 95))}


def measure_throughput(model, images, device, batch: int) -> dict:
    """Batched throughput, which is what whole-volume screening actually uses."""
    import torch

    chunks = [images[i:i + batch] for i in range(0, len(images), batch)]
    model.predict(chunks[0], imgsz=cfg.IMG_SIZE, conf=cfg.INFER_CONF,
                  max_det=cfg.INFER_MAX_DET, device=device, verbose=False)
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    t0 = time.perf_counter()
    n = 0
    for ch in chunks:
        model.predict(ch, imgsz=cfg.IMG_SIZE, conf=cfg.INFER_CONF,
                      max_det=cfg.INFER_MAX_DET, device=device, verbose=False)
        n += len(ch)
    if torch.cuda.is_available():
        torch.cuda.synchronize()
    dt = time.perf_counter() - t0
    return {"batch": batch, "slices": n, "seconds": round(dt, 3),
            "slices_per_second": round(n / dt, 1),
            "seconds_per_scan_256_slices": round(256 * dt / n, 2)}


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=200, help="slices used for timing")
    ap.add_argument("--weights", default=None)
    ap.add_argument("--models", nargs="*", default=["yolo11n.pt", "yolo26n.pt"])
    args = ap.parse_args()

    import torch
    import ultralytics
    from ultralytics import YOLO
    # model_info() returns early (bare `return`) when verbose=False, so query the
    # underlying helpers directly rather than unpacking a None.
    from ultralytics.utils.torch_utils import get_flops, get_num_params

    # Real CT slices, preloaded so disk I/O never enters the timing.
    uid = all_uids()[0]
    vol, _ = read_volume(uid)
    vol_u8 = normalize_hu(vol)
    zs = np.linspace(0, vol_u8.shape[0] - 1, min(args.n, vol_u8.shape[0]), dtype=int)
    images = [np.stack([vol_u8[max(z - 1, 0)], vol_u8[z],
                        vol_u8[min(z + 1, vol_u8.shape[0] - 1)]], axis=-1) for z in zs]
    while len(images) < args.n:
        images += images[:args.n - len(images)]

    device = cfg.HW.device
    report = {
        "platform": {
            "os": platform.platform(),
            "cpu": platform.processor(),
            "logical_cpus": __import__("os").cpu_count(),
            "gpu": torch.cuda.get_device_name(0) if torch.cuda.is_available() else "cpu",
            "gpu_total_memory_GB": round(torch.cuda.get_device_properties(0).total_memory / 1e9, 2)
            if torch.cuda.is_available() else None,
            "torch": torch.__version__,
            "ultralytics": ultralytics.__version__,
            "python": platform.python_version(),
        },
        "settings": {"imgsz": cfg.IMG_SIZE, "conf": cfg.INFER_CONF,
                     "max_det": cfg.INFER_MAX_DET, "amp": cfg.HW.amp},
        "models": {},
    }

    targets = [(Path(args.weights).stem, args.weights)] if args.weights else [
        (m.replace(".pt", ""), str(cfg.WEIGHTS_CACHE / m)) for m in args.models]

    for name, path in targets:
        model = YOLO(path)
        n_params = get_num_params(model.model)
        flops = get_flops(model.model, imgsz=cfg.IMG_SIZE)
        entry = {
            # Recorded by name: an absolute path from the machine that produced
            # the profile is of no use to a reader and leaks a local layout.
            "weights": Path(path).name,
            "size_MB": round(Path(path).stat().st_size / 1e6, 2),
            "parameters": int(n_params),
            "gflops_at_imgsz": round(float(flops), 2),
            "end2end_nms_free": bool(getattr(model.model, "end2end", False)),
            "latency_batch1": measure_latency(model, images, device),
        }
        for b in (16, 32, 64):
            entry[f"throughput_batch{b}"] = measure_throughput(model, images, device, b)
        report["models"][name] = entry
        lat = entry["latency_batch1"]
        print(f"{name:10s} params={entry['parameters']:,} GFLOPs={entry['gflops_at_imgsz']} "
              f"latency={lat['mean_ms']:.2f}+/-{lat['sd_ms']:.2f} ms "
              f"throughput@64={entry['throughput_batch64']['slices_per_second']} slice/s")

    out_json = cfg.RESULTS_DIR / "model_profile.json"
    out_json.write_text(json.dumps(report, indent=1), encoding="utf-8")

    lines = ["# Model and hardware profile", "",
             f"- Platform: {report['platform']['os']}",
             f"- GPU: {report['platform']['gpu']} "
             f"({report['platform']['gpu_total_memory_GB']} GB)",
             f"- torch {report['platform']['torch']}, "
             f"ultralytics {report['platform']['ultralytics']}, "
             f"Python {report['platform']['python']}",
             f"- Input size: {cfg.IMG_SIZE} x {cfg.IMG_SIZE} (native LUNA16 resolution)", "",
             "| Model | Params | GFLOPs | Size (MB) | Latency b=1 (ms) | Throughput b=64 (slice/s) | s / 256-slice scan |",
             "|---|---:|---:|---:|---:|---:|---:|"]
    for name, e in report["models"].items():
        lines.append(
            f"| {name} | {e['parameters']:,} | {e['gflops_at_imgsz']} | {e['size_MB']} | "
            f"{e['latency_batch1']['mean_ms']:.2f} +/- {e['latency_batch1']['sd_ms']:.2f} | "
            f"{e['throughput_batch64']['slices_per_second']} | "
            f"{e['throughput_batch64']['seconds_per_scan_256_slices']} |")
    lines += ["", f"Latency measured over {args.n} consecutive axial slices of a real LUNA16 "
                  "scan preloaded into RAM, batch size 1, timing includes Ultralytics "
                  "pre-processing, forward pass and post-processing, with CUDA "
                  "synchronisation around every call."]
    out_md = cfg.RESULTS_DIR / "model_profile.md"
    out_md.write_text("\n".join(lines), encoding="utf-8")
    print(f"\nWrote {out_json}\n      {out_md}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())