from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

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


def moving_pixel_metrics(prob_move: np.ndarray, label: np.ndarray, threshold: float = 0.5) -> dict[str, float]:
    pred = np.asarray(prob_move) > float(threshold)
    gt = np.asarray(label) == 1

    tp = float(np.logical_and(pred, gt).sum())
    fp = float(np.logical_and(pred, ~gt).sum())
    fn = float(np.logical_and(~pred, gt).sum())

    precision = tp / (tp + fp + 1e-9)
    recall = tp / (tp + fn + 1e-9)
    f1 = 2 * precision * recall / (precision + recall + 1e-9)
    iou = tp / (tp + fp + fn + 1e-9)

    return {
        "moving_precision": precision,
        "moving_recall": recall,
        "moving_f1": f1,
        "moving_iou": iou,
    }


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Evaluate ONNX model on a manifest")
    ap.add_argument("--manifest", required=True)
    ap.add_argument("--onnx", required=True)
    ap.add_argument("--threshold", type=float, default=0.5)
    ap.add_argument("--out", default=None)
    return ap.parse_args()


def main() -> None:
    args = parse_args()
    manifest_path = Path(args.manifest)
    rows = read_jsonl(manifest_path)
    base_dir = manifest_path.resolve().parent

    all_metrics = []
    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")
        m = moving_pixel_metrics(prob_move, lbl, threshold=args.threshold)
        all_metrics.append(m)

    keys = sorted(all_metrics[0].keys()) if all_metrics else []
    summary = {k: float(np.mean([m[k] for m in all_metrics])) for k in keys}
    payload = {
        "n": len(all_metrics),
        "threshold": args.threshold,
        "summary": summary,
    }

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


if __name__ == "__main__":
    main()
