#!/usr/bin/env python3
"""Train one ledger run end to end: train → registry → gold measure → gate → promote.

    python tools/ledger.py new --dataset dataset/2026-10-02_alc --tag alc --base yolo26n-seg.pt
    python tools/train_run.py 2026-10-02_alc --device 0 2>&1 | tee runs/2026-10-02_alc/train.log

Every retraining starts from the base weights on the whole dataset (the cumulative
labelled pool), so runs stay comparable. The gate compares the gold measure with the
deployed model's own gold report (``models/registry/<active>/gold_report.json``):
a pass needs no drop beyond 0.005 in overall F1 and in the mask score (mean mask IoU
over all gold objects, unmatched = 0), an improvement in at least one of them, and no
drop beyond 0.02 in F1 for loop_all or pin. On a pass the
run is promoted on its own: ``models/active/production.json`` is rewritten. The
first run, with no baseline, only records its report as the future baseline.
"""

from __future__ import annotations

import argparse
import csv
import json
import shutil
import socket
import sys
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, Optional

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

from tools import ledger  # noqa: E402

# Recipe since 2026-09-29: no horizontal flip (the camera never shows the pin on the left) and
# full-resolution masks (mask_ratio 1), the combination that kept run A's detection with the
# tightest masks; before that fliplr=0.5 and mask_ratio=4, as in the legacy train_yolo_seg.py.
AUGMENT = dict(scale=0.7, translate=0.15, mosaic=0.5, close_mosaic=10, fliplr=0.0, flipud=0.3,
               shear=5.0, perspective=0.0, mask_ratio=1)
GATE_TOL = 0.005
CLASS_TOL = 0.02
GUARDED_CLASSES = ("loop_all", "pin")


def dataset_dir(run: Dict[str, Any]) -> Path:
    name = str(run["dataset"]).split("/", 1)[-1]
    d = ROOT / "data" / "datasets" / name
    if not (d / "dataset.yaml").exists():
        raise SystemExit(f"dataset {run['dataset']} not found at {d} (git checkout the tag, git lfs pull)")
    return d


def resolved_yaml(run_dir: Path, ds: Path) -> Path:
    text = (ds / "dataset.yaml").read_text(encoding="utf-8").replace("path: .", f"path: {ds.resolve()}", 1)
    out = run_dir / "dataset.resolved.yaml"
    out.write_text(text, encoding="utf-8")
    return out


def base_weights_path(spec: str) -> str:
    p = Path(spec)
    if p.exists():
        return str(p)
    cand = ROOT / "models" / "base" / spec
    return str(cand) if cand.exists() else spec  # else an Ultralytics model name, downloaded


def results_summary(train_dir: Path) -> Dict[str, Any]:
    csv_path = train_dir / "results.csv"
    if not csv_path.exists():
        return {}
    rows = list(csv.DictReader(csv_path.open(encoding="utf-8")))
    if not rows:
        return {}
    last = {k.strip(): v for k, v in rows[-1].items()}
    out: Dict[str, Any] = {"epochs_run": len(rows)}
    for k, v in last.items():
        if "metrics/" in k:
            try:
                out[k.replace("metrics/", "")] = float(v)
            except ValueError:
                pass
    return out


def measure(weights: Path, run_dir: Path, names: list, conf: float = 0.25) -> Dict[str, Any]:
    from src.eval.evaluate import evaluate_dirs
    from src.eval.yolo_labels import PolygonInstance, write_label_file
    from ultralytics import YOLO

    gold = ROOT / "data" / "gold"
    img_dir, gold_labels = gold / "images", gold / "labels"
    if not img_dir.exists() or not gold_labels.exists():
        raise SystemExit("no gold set under data/gold (import it with tools/import_studio_export.py --gold)")
    pred_dir = run_dir / "gold_predictions"
    pred_dir.mkdir(parents=True, exist_ok=True)
    model = YOLO(str(weights))
    for img_path in sorted(p for p in img_dir.iterdir() if p.suffix.lower() in (".jpg", ".jpeg", ".png")):
        res = model.predict(str(img_path), conf=conf, verbose=False)[0]
        insts = []
        if res.masks is not None:
            h, w = res.orig_shape
            for poly, cls in zip(res.masks.xy, res.boxes.cls.tolist()):
                pts = [(float(x) / w, float(y) / h) for x, y in poly]
                if len(pts) >= 3:
                    insts.append(PolygonInstance(int(cls), pts))
        write_label_file(pred_dir / f"{img_path.stem}.txt", insts)
    return evaluate_dirs(pred_dir, gold_labels, names=names)


def mask_score(overall: Dict[str, Any]) -> float:
    """Mean mask IoU over every gold object, an unmatched object counting as 0.

    ``mean_mask_iou`` alone is averaged over matched objects only, so a model that
    finds only the easy objects scores higher than one that also finds the hard ones
    with looser masks; run 1 on 2026-09-28 lost the gate exactly that way. Weighting
    by recall removes that bias.
    """
    tp, fn = float(overall.get("tp", 0)), float(overall.get("fn", 0))
    return float(overall.get("mean_mask_iou", 0.0)) * (tp / (tp + fn) if tp + fn else 0.0)


def gate(report: Dict[str, Any], baseline: Optional[Path]) -> Optional[Dict[str, Any]]:
    if not baseline or not Path(baseline).exists():
        return None
    base = json.loads(Path(baseline).read_text(encoding="utf-8"))
    f1_delta = float(report["overall"]["f1"]) - float(base["overall"]["f1"])
    iou_delta = mask_score(report["overall"]) - mask_score(base["overall"])
    per_class = report.get("per_class") or report.get("classes") or {}
    base_classes = base.get("per_class") or base.get("classes") or {}
    class_deltas = {}
    for c in GUARDED_CLASSES:
        a, b = per_class.get(c) or {}, base_classes.get(c) or {}
        if "f1" in a and "f1" in b:
            class_deltas[c] = float(a["f1"]) - float(b["f1"])
    ok = (f1_delta >= -GATE_TOL and iou_delta >= -GATE_TOL and (f1_delta > 0 or iou_delta > 0)
          and all(d >= -CLASS_TOL for d in class_deltas.values()))
    return {"baseline": str(baseline), "f1_delta": round(f1_delta, 4), "mask_iou_delta": round(iou_delta, 4),
            "mask_score": round(mask_score(report["overall"]), 4), "mask_score_baseline": round(mask_score(base["overall"]), 4),
            "matched_only_iou_delta": round(float(report["overall"]["mean_mask_iou"]) - float(base["overall"]["mean_mask_iou"]), 4),
            "class_f1_delta": {k: round(v, 4) for k, v in class_deltas.items()}, "pass": bool(ok)}


def active_version() -> Optional[str]:
    p = ROOT / "models" / "active" / "production.json"
    if not p.exists():
        return None
    return json.loads(p.read_text(encoding="utf-8")).get("version")


def promote(version: str, run_id: str, g: Dict[str, Any]) -> Path:
    p = ROOT / "models" / "active" / "production.json"
    p.parent.mkdir(parents=True, exist_ok=True)
    p.write_text(json.dumps({
        "version": version,
        "export": "weights.pt",   # engine built on the target GPU at deploy time
        "promoted_on": datetime.now().strftime("%Y-%m-%d"),
        "promoted_by": "gate",
        "run": run_id,
        "gate": g,
    }, indent=2) + "\n", encoding="utf-8")
    return p


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("run_id")
    ap.add_argument("--epochs", type=int, default=400)
    ap.add_argument("--patience", type=int, default=50)
    ap.add_argument("--imgsz", type=int, default=640)
    ap.add_argument("--batch", type=int, default=16)
    ap.add_argument("--device", default="0")
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--conf", type=float, default=0.25, help="confidence used for the gold measure")
    ap.add_argument("--baseline", default=None, help="gold report to gate against (default: the active model's)")
    ap.add_argument("--no-promote", action="store_true", help="gate, but leave production.json alone")
    ap.add_argument("--skip-train", action="store_true", help="reuse runs/<id>/ultralytics/train/weights/best.pt")
    ap.add_argument("--set", nargs="*", default=[], metavar="KEY=VALUE",
                    help="override any Ultralytics train argument, e.g. fliplr=0 mask_ratio=1 overlap_mask=true; recorded in the ledger")
    args = ap.parse_args(argv)

    run = ledger.load(args.run_id)
    run_dir = ledger.run_path(args.run_id).parent
    ds = dataset_dir(run)
    data_yaml = resolved_yaml(run_dir, ds)
    names = [v for _, v in sorted(json.loads((ds / "source.json").read_text())["classes"].items(), key=lambda kv: int(kv[0]))] \
        if (ds / "source.json").exists() else ["loop_all", "pin", "crystal", "loop_face", "ice", "needle"]
    overrides: Dict[str, Any] = {}
    for kv in args.set or []:
        if "=" not in kv:
            raise SystemExit(f"--set expects KEY=VALUE, got {kv!r}")
        k, v = kv.split("=", 1)
        try:
            overrides[k] = json.loads(v.lower() if v.lower() in ("true", "false") else v)
        except json.JSONDecodeError:
            overrides[k] = v
    hyper = {"epochs": args.epochs, "patience": args.patience, "imgsz": args.imgsz, "batch": args.batch,
             "device": args.device, "overlap_mask": False, **AUGMENT, **overrides}
    ledger.set_status(args.run_id, "training", note="train_run.py", kv={
        "hyperparams": hyper, "host": socket.gethostname(), "started": datetime.now().isoformat(timespec="seconds")})

    best = run_dir / "ultralytics" / "train" / "weights" / "best.pt"
    if not args.skip_train:
        from ultralytics import YOLO
        model = YOLO(base_weights_path(run["base_weights"]))
        train_kwargs = dict(epochs=args.epochs, patience=args.patience, imgsz=args.imgsz, batch=args.batch,
                            device=args.device, workers=args.workers, overlap_mask=False, amp=False, **AUGMENT)
        train_kwargs.update(overrides)   # --set wins over the recipe
        model.train(data=str(data_yaml), project=str(run_dir / "ultralytics"), name="train", exist_ok=True, **train_kwargs)
    if not best.exists():
        ledger.set_status(args.run_id, "aborted", note=f"no best.pt at {best}")
        return 1

    family = Path(run["base_weights"]).stem.split("-")[0]
    tag = args.run_id.split("_", 1)[-1]
    version = f"{family}-seg-{tag}-{datetime.now().strftime('%Y-%m-%d')}"
    reg = ROOT / "models" / "registry" / version
    reg.mkdir(parents=True, exist_ok=True)
    shutil.copy2(best, reg / "weights.pt")

    report = measure(reg / "weights.pt", run_dir, names, conf=args.conf)
    baseline = Path(args.baseline) if args.baseline else None
    if baseline is None:
        av = active_version()
        if av and (ROOT / "models" / "registry" / av / "gold_report.json").exists():
            baseline = ROOT / "models" / "registry" / av / "gold_report.json"
    g = gate(report, baseline)
    report["gate"] = g
    (reg / "gold_report.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    metrics = {"overall": report["overall"], "gate": g, "training": results_summary(run_dir / "ultralytics" / "train"),
               "version": version, "measured_at": datetime.now().isoformat(timespec="seconds")}
    (run_dir / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n", encoding="utf-8")
    (reg / "metadata.json").write_text(json.dumps({
        "model_id": version, "family": family, "task": "segmentation",
        "training_date": datetime.now().strftime("%Y-%m-%d"), "run": args.run_id,
        "dataset": run["dataset"], "base_weights": run["base_weights"], "hyperparams": hyper,
        "export_formats": ["pt"], "metrics": metrics["overall"], "gate": g,
    }, indent=2) + "\n", encoding="utf-8")
    ledger.record_metrics(args.run_id, run_dir / "metrics.json")
    ledger.set_status(args.run_id, "measured", kv={"version": version})

    if g is None:
        print(f"measured; no baseline, {version} is the baseline for the next run")
        return 0
    if g["pass"]:
        ledger.set_status(args.run_id, "gate_passed", note=json.dumps(g))
        if not args.no_promote:
            promote(version, args.run_id, g)
            ledger.set_status(args.run_id, "promoted", note="by the gate")
            print(f"PROMOTED {version}: {g}")
        return 0
    ledger.set_status(args.run_id, "gate_failed", note=json.dumps(g))
    print(f"gate failed: {g}")
    return 2


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