#!/usr/bin/env python3
"""Cross-dataset evaluation of the PHOENIX ground-sensor models on PyroNear2025.

WHY THIS EXISTS
---------------
Our headline 98.0% recall is measured on our own Sicilian ground set. That answers
"is it skilled here?", not "does it generalise?". An independent reviewer reaches for
PyroNear2025 (pyronear/pyro-sdis) instead. Both numbers belong in the record, side by
side, with the difference stated plainly.

An earlier run of this evaluation produced a figure that is quoted in the regulatory
register. Its result file lived in /tmp and did not survive the 2026-09-06 02:11
reboot, so the quoted number is currently unreproducible. This run re-derives it and
writes the output somewhere durable.

TWO TRAPS THIS SCRIPT EXISTS TO AVOID
-------------------------------------
1. CLASS INDEX. The pyro-sdis parquet labels smoke as class 1. Our models emit smoke
   as class 0. The materialised labels under yolo_val/ have ALREADY been remapped
   (verified: 3942 annotations, all class 0) and the yaml declares both names so the
   index space matches the model. This script asserts that rather than assuming it —
   a silent disagreement here reports 0.000 for everything and looks like a bad model.

2. STALE CACHE. ultralytics writes labels/val.cache and reuses it without checking
   whether the labels changed underneath. A cache built before the remap made the
   first two attempts report all zeros. The cache is deleted before every run.

Both failures are silent and both look exactly like "the model is worthless", which is
why they are checked rather than trusted.
"""
import json
import os
import shutil
import subprocess
import sys
import time
from collections import Counter
from pathlib import Path

# PYRO_DATA_DIR points at a local-NVMe copy when one has been staged. Reading the
# split straight off /media ran at 11.9 MB/s and left the GPU idle between batches.
DATA = Path(os.environ.get("PYRO_DATA_DIR",
                           "/media/mark/AI_DGX1/Wildfire_AI/pyro_sdis/yolo_val"))
YAML = DATA / "pyro_sdis.yaml"
WEIGHT_ROOT = Path("/media/mark/AI_DGX1/Wildfire_AI/phoenix_ground_weights")
OUT_DIR = WEIGHT_ROOT / "pyro_eval"          # durable: NOT /tmp, survives a reboot
IMGSZ = 1024
SMOKE_CLASS = 0

MODELS = {
    "v15_SELFDISTILL": WEIGHT_ROOT / "yolov8s_v15_SELFDISTILL_imgsz1024" / "best.pt",
    "v20_HARDNEG": WEIGHT_ROOT / "yolov8s_v20_HARDNEG_imgsz1024" / "best.pt",
}

# ── abort conditions ─────────────────────────────────────────────────────────
# load_guard.sh owns thermal and contention. This script owns the two things
# load_guard cannot know about: PHOENIX must be healthy, and the dataset must be
# in the state the evaluation assumes.


def phoenix_healthy() -> bool:
    r = subprocess.run(["systemctl", "is-active", "adr-wildfire-solution.service"],
                       capture_output=True, text=True)
    return r.stdout.strip() == "active"


def check_dataset() -> dict:
    """Prove the label/model class spaces agree BEFORE spending GPU time."""
    counts: Counter = Counter()
    empty = 0
    label_files = sorted((DATA / "labels" / "val").glob("*.txt"))
    for p in label_files:
        rows = [ln for ln in p.read_text().splitlines() if ln.strip()]
        if not rows:
            empty += 1
        for ln in rows:
            counts[int(ln.split()[0])] += 1
    n_images = len(list((DATA / "images" / "val").glob("*.jpg")))

    if not counts:
        sys.exit("ABORT: no annotations found — dataset not materialised")
    if set(counts) != {SMOKE_CLASS}:
        sys.exit(f"ABORT: labels use classes {sorted(counts)}, model emits "
                 f"{SMOKE_CLASS}. Remap was not applied — a run now would report "
                 f"all zeros and look like a model failure.")
    return {"images": n_images, "label_files": len(label_files),
            "annotations": counts[SMOKE_CLASS], "empty_label_files": empty}


def drop_caches() -> list:
    """ultralytics reuses a cache built against the pre-remap labels."""
    killed = []
    for c in DATA.glob("labels/*.cache"):
        c.unlink()
        killed.append(str(c))
    return killed


def evaluate(name: str, weights: Path) -> dict:
    from ultralytics import YOLO

    drop_caches()
    model = YOLO(str(weights))
    t0 = time.time()
    m = model.val(data=str(YAML), imgsz=IMGSZ, split="val", batch=8,
                  conf=0.001, iou=0.6, plots=False, verbose=False,
                  project=str(OUT_DIR / "runs"), name=name, exist_ok=True)
    secs = time.time() - t0

    b = m.box
    # class 0 is smoke; index into the per-class arrays by position in ap_class_index
    idx = list(b.ap_class_index).index(SMOKE_CLASS) if SMOKE_CLASS in b.ap_class_index else None
    p = float(b.p[idx]) if idx is not None else 0.0
    r = float(b.r[idx]) if idx is not None else 0.0
    f1 = (2 * p * r / (p + r)) if (p + r) else 0.0

    return {
        "model": name,
        "weights": str(weights),
        "precision": round(p, 4),
        "recall": round(r, 4),
        "f1": round(f1, 4),
        "ap50": round(float(b.ap50[idx]) if idx is not None else 0.0, 4),
        "ap50_95": round(float(b.ap[idx]) if idx is not None else 0.0, 4),
        "map50": round(float(b.map50), 4),
        "imgsz": IMGSZ,
        "seconds": round(secs, 1),
    }


def main() -> int:
    if not phoenix_healthy():
        print("ABORT: adr-wildfire-solution.service is not active")
        return 1

    OUT_DIR.mkdir(parents=True, exist_ok=True)
    ds = check_dataset()
    print(f"dataset OK: {ds['images']} images, {ds['annotations']} smoke boxes, "
          f"{ds['empty_label_files']} background frames")

    results = {}
    for name, w in MODELS.items():
        out = OUT_DIR / f"{name}.json"
        if out.exists():                      # idempotent: resume after a SIGTERM
            print(f"{name}: already done, skipping")
            results[name] = json.loads(out.read_text())
            continue
        if not w.exists():
            print(f"{name}: weights missing at {w} — SKIPPED")
            continue
        if not phoenix_healthy():             # re-check between models
            print("ABORT: PHOENIX went unhealthy mid-run")
            return 1
        print(f"{name}: evaluating {w}")
        res = evaluate(name, w)
        res["dataset"] = ds
        res["dataset_name"] = "pyronear/pyro-sdis val (PyroNear2025)"
        res["run_at"] = time.strftime("%Y-%m-%dT%H:%M:%S%z")
        out.write_text(json.dumps(res, indent=2))
        results[name] = res
        print(f"  P {res['precision']}  R {res['recall']}  F1 {res['f1']}  "
              f"AP50 {res['ap50']}  ({res['seconds']}s)")

    summary = OUT_DIR / "summary.json"
    summary.write_text(json.dumps(results, indent=2))
    print(f"\nwrote {summary}")
    for n, r in results.items():
        print(f"  {n:20s} F1 {r['f1']:.3f}  AP50 {r['ap50']:.3f}  R {r['recall']:.3f}")
    return 0


if __name__ == "__main__":
    sys.exit(main())