from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from .eval import moving_pixel_metrics
from .infer_onnx import infer_prob_move_onnx
from .io import load_image, percentile_normalize, read_jsonl, resolve_manifest_path


def _parse_threshold_list(raw: str) -> list[float]:
    values: list[float] = []
    for item in str(raw).split(","):
        txt = item.strip()
        if not txt:
            continue
        v = float(txt)
        if not (0.0 <= v <= 1.0):
            raise ValueError(f"Threshold must be in [0, 1], got {v}")
        values.append(float(v))
    if not values:
        raise ValueError("No valid thresholds were provided in --thresholds")
    # Stable unique + sorted
    uniq = sorted(set(round(v, 10) for v in values))
    return [float(v) for v in uniq]


def _build_thresholds(start: float, stop: float, step: float) -> list[float]:
    if step <= 0:
        raise ValueError("--step must be > 0")
    if stop < start:
        raise ValueError("--stop must be >= --start")
    if start < 0.0 or stop > 1.0:
        raise ValueError("--start/--stop must be in [0, 1]")

    thresholds: list[float] = []
    i = 0
    while True:
        val = start + i * step
        if val > stop + 1e-12:
            break
        thresholds.append(float(round(val, 10)))
        i += 1

    if not thresholds:
        thresholds = [float(round(start, 10))]
    return thresholds


def _summarize(metrics_rows: list[dict[str, float]]) -> dict[str, float]:
    if not metrics_rows:
        return {
            "moving_precision": 0.0,
            "moving_recall": 0.0,
            "moving_f1": 0.0,
            "moving_iou": 0.0,
        }
    keys = sorted(metrics_rows[0].keys())
    return {k: float(np.mean([row[k] for row in metrics_rows])) for k in keys}


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Sweep thresholds for ONNX moving-class metrics")
    ap.add_argument("--manifest", required=True, help="JSONL manifest with image/label entries")
    ap.add_argument("--onnx", required=True, help="Path to ONNX model")
    ap.add_argument("--thresholds", default=None, help="Comma-separated thresholds (overrides start/stop/step)")
    ap.add_argument("--start", type=float, default=0.05, help="Threshold sweep start (inclusive)")
    ap.add_argument("--stop", type=float, default=0.95, help="Threshold sweep stop (inclusive)")
    ap.add_argument("--step", type=float, default=0.05, help="Threshold sweep step")
    ap.add_argument(
        "--recall_min_precision",
        type=float,
        default=0.0,
        help="When selecting best-recall threshold, require at least this precision",
    )
    ap.add_argument("--out", default=None, help="Optional output JSON path")
    return ap.parse_args()


def main() -> None:
    args = parse_args()

    if args.thresholds:
        thresholds = _parse_threshold_list(args.thresholds)
    else:
        thresholds = _build_thresholds(float(args.start), float(args.stop), float(args.step))

    recall_min_precision = float(args.recall_min_precision)
    if not (0.0 <= recall_min_precision <= 1.0):
        raise ValueError("--recall_min_precision must be in [0, 1]")

    manifest_path = Path(args.manifest)
    rows = read_jsonl(manifest_path)
    base_dir = manifest_path.resolve().parent

    # Run ONNX once per sample, then sweep thresholds on cached probability maps.
    probs: list[np.ndarray] = []
    labels: list[np.ndarray] = []
    for row in rows:
        img_path = resolve_manifest_path(base_dir, row["image"])
        lbl_path = resolve_manifest_path(base_dir, row["label"])

        img = percentile_normalize(load_image(img_path))
        lbl = load_image(lbl_path).astype(np.int64)
        prob_move = infer_prob_move_onnx(img, args.onnx, time_axis="y")

        probs.append(np.asarray(prob_move, dtype=np.float32))
        labels.append(np.asarray(lbl, dtype=np.int64))

    sweep_rows = []
    for thr in thresholds:
        per_sample = [
            moving_pixel_metrics(prob, lbl, threshold=float(thr))
            for prob, lbl in zip(probs, labels)
        ]
        summary = _summarize(per_sample)
        sweep_rows.append(
            {
                "threshold": float(thr),
                "summary": summary,
            }
        )

    if sweep_rows:
        best_f1 = max(
            sweep_rows,
            key=lambda r: (
                float(r["summary"]["moving_f1"]),
                float(r["summary"]["moving_recall"]),
                float(r["summary"]["moving_precision"]),
            ),
        )
        recall_candidates = [
            r for r in sweep_rows
            if float(r["summary"]["moving_precision"]) >= recall_min_precision
        ]
        if not recall_candidates:
            recall_candidates = sweep_rows
        best_recall = max(
            recall_candidates,
            key=lambda r: (
                float(r["summary"]["moving_recall"]),
                float(r["summary"]["moving_f1"]),
                float(r["summary"]["moving_precision"]),
            ),
        )
    else:
        best_f1 = {"threshold": None, "summary": _summarize([])}
        best_recall = {"threshold": None, "summary": _summarize([])}

    payload = {
        "n": len(probs),
        "thresholds": thresholds,
        "rows": sweep_rows,
        "best": {
            "by_f1": best_f1,
            "by_recall": best_recall,
            "recall_min_precision": recall_min_precision,
        },
    }

    print(f"Loaded {len(probs)} sample(s). Swept {len(thresholds)} threshold(s).")
    for row in sweep_rows:
        s = row["summary"]
        print(
            "thr={thr:.3f} f1={f1:.4f} rec={rec:.4f} prec={prec:.4f} iou={iou:.4f}".format(
                thr=float(row["threshold"]),
                f1=float(s["moving_f1"]),
                rec=float(s["moving_recall"]),
                prec=float(s["moving_precision"]),
                iou=float(s["moving_iou"]),
            )
        )

    bf = best_f1["summary"]
    br = best_recall["summary"]
    print(
        "Best by F1: thr={thr:.3f} f1={f1:.4f} rec={rec:.4f} prec={prec:.4f}".format(
            thr=float(best_f1["threshold"]),
            f1=float(bf["moving_f1"]),
            rec=float(bf["moving_recall"]),
            prec=float(bf["moving_precision"]),
        )
    )
    print(
        "Best by recall: thr={thr:.3f} rec={rec:.4f} f1={f1:.4f} prec={prec:.4f}".format(
            thr=float(best_recall["threshold"]),
            rec=float(br["moving_recall"]),
            f1=float(br["moving_f1"]),
            prec=float(br["moving_precision"]),
        )
    )

    if args.out:
        Path(args.out).write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")


if __name__ == "__main__":
    main()
