from __future__ import annotations

import argparse
import json
import time
from pathlib import Path
from typing import Any

import numpy as np

from .infer_onnx import infer_prob_move_and_embeddings_onnx
from .io import load_image, percentile_normalize, read_jsonl, resolve_manifest_path
from .postprocess import postprocess_prob_to_tracks


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")
    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]")
    out: list[float] = []
    i = 0
    while True:
        val = start + i * step
        if val > stop + 1e-12:
            break
        out.append(float(round(val, 10)))
        i += 1
    if not out:
        out = [float(round(start, 10))]
    return out


def _parse_track_points(track: dict[str, Any]) -> tuple[list[tuple[int, float]], int | None]:
    max_gap = track.get("max_interp_gap", None)
    if max_gap is None:
        max_gap_i = None
    else:
        try:
            max_gap_i = int(max_gap)
        except Exception:
            max_gap_i = None

    pts: list[tuple[int, float]] = []
    raw_path = track.get("path_tx")
    if isinstance(raw_path, list):
        for node in raw_path:
            if not isinstance(node, (list, tuple)) or len(node) < 2:
                continue
            try:
                t = int(round(float(node[0])))
                x = float(node[1])
            except Exception:
                continue
            if not np.isfinite(x):
                continue
            pts.append((t, x))

    if not pts:
        raw_anchors = track.get("anchors")
        if isinstance(raw_anchors, list):
            for a in raw_anchors:
                if isinstance(a, dict):
                    try:
                        t = int(round(float(a.get("t"))))
                        x = float(a.get("x"))
                    except Exception:
                        continue
                elif isinstance(a, (list, tuple)) and len(a) >= 2:
                    try:
                        t = int(round(float(a[0])))
                        x = float(a[1])
                    except Exception:
                        continue
                else:
                    continue
                if not np.isfinite(x):
                    continue
                pts.append((t, x))
        if max_gap_i is None:
            # Anchors were historically connected without gap restriction.
            max_gap_i = 0

    return pts, max_gap_i


def _dedupe_points(points: list[tuple[int, float]]) -> list[tuple[int, float]]:
    if not points:
        return []
    points = sorted(points, key=lambda z: z[0])
    out: list[tuple[int, float]] = []
    cur_t = points[0][0]
    xs: list[float] = []
    for t, x in points:
        if t != cur_t:
            out.append((int(cur_t), float(np.mean(xs))))
            cur_t = t
            xs = [x]
        else:
            xs.append(x)
    out.append((int(cur_t), float(np.mean(xs))))
    return out


def _dense_track(points: list[tuple[int, float]], max_interp_gap: int | None) -> dict[int, float]:
    pts = _dedupe_points(points)
    if len(pts) < 2:
        return {}

    by_t: dict[int, list[float]] = {}
    for (t0, x0), (t1, x1) in zip(pts[:-1], pts[1:]):
        by_t.setdefault(int(t0), []).append(float(x0))
        if t1 <= t0:
            continue
        dt = int(t1 - t0)
        # Same semantics as tracy_to_real rasterization:
        # if max_interp_gap > 0 and gap > max_interp_gap, do not bridge.
        if max_interp_gap is not None and max_interp_gap > 0 and dt > int(max_interp_gap):
            continue
        for t in range(int(t0) + 1, int(t1)):
            a = float((t - t0) / max(1, dt))
            x = float((1.0 - a) * x0 + a * x1)
            by_t.setdefault(int(t), []).append(float(x))
    by_t.setdefault(int(pts[-1][0]), []).append(float(pts[-1][1]))

    out: dict[int, float] = {}
    for t, xs in by_t.items():
        out[int(t)] = float(np.mean(xs))
    return out


def _build_gt_tracks(meta: dict[str, Any]) -> tuple[list[dict[int, float]], bool]:
    tracks = meta.get("tracks", [])
    if not isinstance(tracks, list):
        return [], False

    out: list[dict[int, float]] = []
    missing_identity_paths = False
    for tr in tracks:
        if not isinstance(tr, dict):
            continue
        cls = str(tr.get("class", "")).strip().lower()
        # Align identity evaluation with moving-class inference:
        # ignore static GT tracks when class labels are available.
        if cls and cls == "static":
            continue
        pts, max_gap = _parse_track_points(tr)
        if not pts:
            # Legacy real metadata had only trajectory/start/end stats.
            # If there are declared tracks but no points, identity eval is not possible.
            missing_identity_paths = True
            continue
        dense = _dense_track(pts, max_interp_gap=max_gap)
        if len(dense) >= 2:
            out.append(dense)

    return out, missing_identity_paths


def _build_pred_tracks(pred_tracks: list[dict[str, Any]]) -> list[dict[int, float]]:
    out: list[dict[int, float]] = []
    for tr in pred_tracks:
        if not isinstance(tr, dict):
            continue
        points: list[tuple[int, float]] = []
        raw_path = tr.get("path")
        if isinstance(raw_path, list):
            for node in raw_path:
                if not isinstance(node, (list, tuple)) or len(node) < 2:
                    continue
                try:
                    t = int(round(float(node[0])))
                    x = float(node[1])
                except Exception:
                    continue
                if not np.isfinite(x):
                    continue
                points.append((t, x))
        if not points:
            anchors = tr.get("anchors", [])
            if isinstance(anchors, list):
                for a in anchors:
                    if not isinstance(a, dict):
                        continue
                    try:
                        t = int(round(float(a.get("t"))))
                        x = float(a.get("x"))
                    except Exception:
                        continue
                    if not np.isfinite(x):
                        continue
                    points.append((t, x))
        dense = _dense_track(points, max_interp_gap=0)
        if len(dense) >= 2:
            out.append(dense)
    return out


def _pair_matches(a: dict[int, float], b: dict[int, float], tol_px: float) -> int:
    if not a or not b:
        return 0
    keys = set(a.keys()).intersection(b.keys())
    if not keys:
        return 0
    tol = float(max(0.0, tol_px))
    return int(sum(1 for t in keys if abs(float(a[t]) - float(b[t])) <= tol))


def _assignment_max(weight: np.ndarray) -> list[tuple[int, int]]:
    if weight.size == 0:
        return []
    try:
        from scipy.optimize import linear_sum_assignment

        rr, cc = linear_sum_assignment(-weight)
        return [(int(r), int(c)) for r, c in zip(rr, cc) if weight[int(r), int(c)] > 0]
    except Exception:
        # Fallback greedy assignment.
        pairs: list[tuple[int, int, float]] = []
        for i in range(weight.shape[0]):
            for j in range(weight.shape[1]):
                w = float(weight[i, j])
                if w > 0:
                    pairs.append((i, j, w))
        pairs.sort(key=lambda x: x[2], reverse=True)
        out: list[tuple[int, int]] = []
        used_r: set[int] = set()
        used_c: set[int] = set()
        for i, j, _ in pairs:
            if i in used_r or j in used_c:
                continue
            used_r.add(i)
            used_c.add(j)
            out.append((int(i), int(j)))
        return out


def _safe_prf(tp: float, fp: float, fn: float) -> tuple[float, float, float]:
    prec = float(tp / (tp + fp + 1e-9))
    rec = float(tp / (tp + fn + 1e-9))
    f1 = float(2.0 * prec * rec / (prec + rec + 1e-9))
    return prec, rec, f1


def _eval_sample_identity(
    gt_tracks: list[dict[int, float]],
    pred_tracks: list[dict[int, float]],
    *,
    tol_px: float,
    track_iou_thresh: float,
    track_min_match_frames: int,
) -> dict[str, float]:
    n_gt = len(gt_tracks)
    n_pred = len(pred_tracks)
    gt_pts = float(sum(len(tr) for tr in gt_tracks))
    pred_pts = float(sum(len(tr) for tr in pred_tracks))

    match = np.zeros((n_gt, n_pred), dtype=np.float32)
    for i, g in enumerate(gt_tracks):
        for j, p in enumerate(pred_tracks):
            match[i, j] = float(_pair_matches(g, p, tol_px=float(tol_px)))

    pairs = _assignment_max(match)
    idtp = float(sum(match[i, j] for i, j in pairs))
    idfp = float(max(0.0, pred_pts - idtp))
    idfn = float(max(0.0, gt_pts - idtp))

    tp_tracks = 0.0
    for i, j in pairs:
        m = float(match[i, j])
        if m <= 0:
            continue
        g_len = float(max(1, len(gt_tracks[i])))
        p_len = float(max(1, len(pred_tracks[j])))
        iou = float(m / max(1e-9, (g_len + p_len - m)))
        if m >= float(track_min_match_frames) and iou >= float(track_iou_thresh):
            tp_tracks += 1.0

    fp_tracks = float(max(0.0, n_pred - tp_tracks))
    fn_tracks = float(max(0.0, n_gt - tp_tracks))

    return {
        "idtp": idtp,
        "idfp": idfp,
        "idfn": idfn,
        "tp_tracks": tp_tracks,
        "fp_tracks": fp_tracks,
        "fn_tracks": fn_tracks,
        "gt_tracks": float(n_gt),
        "pred_tracks": float(n_pred),
    }


def _summary_from_counts(counts: dict[str, float]) -> dict[str, float]:
    idp, idr, idf1 = _safe_prf(counts["idtp"], counts["idfp"], counts["idfn"])
    trk_p, trk_r, trk_f1 = _safe_prf(counts["tp_tracks"], counts["fp_tracks"], counts["fn_tracks"])
    return {
        "id_precision": float(idp),
        "id_recall": float(idr),
        "id_f1": float(idf1),
        "track_precision": float(trk_p),
        "track_recall": float(trk_r),
        "track_f1": float(trk_f1),
        "idtp": float(counts["idtp"]),
        "idfp": float(counts["idfp"]),
        "idfn": float(counts["idfn"]),
        "tp_tracks": float(counts["tp_tracks"]),
        "fp_tracks": float(counts["fp_tracks"]),
        "fn_tracks": float(counts["fn_tracks"]),
        "gt_tracks": float(counts["gt_tracks"]),
        "pred_tracks": float(counts["pred_tracks"]),
    }


def parse_args() -> argparse.Namespace:
    ap = argparse.ArgumentParser(description="Sweep thresholds for identity-aware track metrics")
    ap.add_argument("--manifest", required=True, help="JSONL manifest with image/label/meta entries")
    ap.add_argument("--onnx", required=True, help="Path to ONNX model")
    ap.add_argument("--max_samples", type=int, default=0, help="Evaluate at most this many manifest rows (0=all)")
    ap.add_argument("--log_every", type=int, default=100, help="Progress print interval for sample loops")
    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 identity precision",
    )
    ap.add_argument("--tol_px", type=float, default=2.5, help="Match tolerance in x (pixels) for per-frame IDTP")
    ap.add_argument("--track_iou_thresh", type=float, default=0.30, help="Track-level match IoU threshold")
    ap.add_argument(
        "--track_min_match_frames",
        type=int,
        default=6,
        help="Track-level match requires at least this many matched frames",
    )

    # Postprocess parameters aligned with Tracy FIND defaults.
    ap.add_argument("--min_len", type=float, default=3.0)
    ap.add_argument("--vmin", type=float, default=0.0)
    ap.add_argument("--simplify_eps", type=float, default=0.9)
    ap.add_argument("--max_anchors", type=int, default=24)
    ap.add_argument("--min_component_size", type=int, default=3)
    ap.add_argument("--min_t_span", type=int, default=6)
    ap.add_argument("--min_mean_prob", type=float, default=0.11)
    ap.add_argument("--max_tracks", type=int, default=100)
    ap.add_argument("--smooth_sigma_t", type=float, default=0.05)
    ap.add_argument("--smooth_sigma_x", type=float, default=0.25)
    ap.add_argument("--merge_gap_t", type=int, default=1)
    ap.add_argument("--merge_max_dx", type=float, default=2.6)
    ap.add_argument("--min_directionality", type=float, default=0.0)
    ap.add_argument("--min_net_velocity", type=float, default=0.0)
    ap.add_argument("--merge_min_bridge_prob", type=float, default=0.30)
    ap.add_argument("--merge_min_bridge_coverage", type=float, default=0.80)
    ap.add_argument(
        "--merge_min_embed_sim",
        type=float,
        default=0.18,
        help="Minimum endpoint embedding cosine similarity for merge (if embeddings are available)",
    )
    ap.add_argument("--out", default=None, help="Optional output JSON path")
    return ap.parse_args()


def main() -> None:
    args = parse_args()
    t0_all = time.time()

    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)
    if int(args.max_samples) > 0:
        rows = rows[: int(args.max_samples)]
    base_dir = manifest_path.resolve().parent
    log_every = max(1, int(args.log_every))

    probs: list[np.ndarray] = []
    emb_maps: list[np.ndarray | None] = []
    gt_tracks_all: list[list[dict[int, float]]] = []
    skipped_no_meta = 0
    skipped_missing_paths = 0

    t0_load = time.time()
    print(f"[info] loading + ONNX inference for {len(rows)} row(s)...", flush=True)
    for i, row in enumerate(rows, start=1):
        meta_rel = row.get("meta")
        if not meta_rel:
            skipped_no_meta += 1
            if (i % log_every) == 0 or i == len(rows):
                print(
                    f"[info] infer progress: {i}/{len(rows)} "
                    f"(kept={len(probs)} no_meta={skipped_no_meta} missing_paths={skipped_missing_paths})",
                    flush=True,
                )
            continue

        img_path = resolve_manifest_path(base_dir, row["image"])
        meta_path = resolve_manifest_path(base_dir, meta_rel)

        img = percentile_normalize(load_image(img_path))
        prob_move, emb_map = infer_prob_move_and_embeddings_onnx(img, args.onnx, time_axis="y")

        try:
            meta = json.loads(Path(meta_path).read_text(encoding="utf-8"))
        except Exception:
            skipped_no_meta += 1
            continue

        gt_tracks, missing_paths = _build_gt_tracks(meta)
        declared_n_tracks = int(meta.get("n_tracks", len(meta.get("tracks", []) or [])))
        if declared_n_tracks > 0 and not gt_tracks and missing_paths:
            skipped_missing_paths += 1
            if (i % log_every) == 0 or i == len(rows):
                print(
                    f"[info] infer progress: {i}/{len(rows)} "
                    f"(kept={len(probs)} no_meta={skipped_no_meta} missing_paths={skipped_missing_paths})",
                    flush=True,
                )
            continue

        probs.append(np.asarray(prob_move, dtype=np.float32))
        emb_maps.append(np.asarray(emb_map, dtype=np.float32) if emb_map is not None else None)
        gt_tracks_all.append(gt_tracks)
        if (i % log_every) == 0 or i == len(rows):
            print(
                f"[info] infer progress: {i}/{len(rows)} "
                f"(kept={len(probs)} no_meta={skipped_no_meta} missing_paths={skipped_missing_paths})",
                flush=True,
            )

    if not probs:
        raise RuntimeError(
            "No evaluable samples found. If using real data converted by older tracy_to_real, "
            "re-convert so each meta track includes path_tx."
        )
    print(f"[info] inference stage done in {(time.time() - t0_load):.1f}s", flush=True)

    post_args = {
        "min_len": float(args.min_len),
        "vmin": float(args.vmin),
        "simplify_eps": float(args.simplify_eps),
        "max_anchors": int(args.max_anchors),
        "min_component_size": int(args.min_component_size),
        "min_t_span": int(args.min_t_span),
        "min_mean_prob": float(args.min_mean_prob),
        "max_tracks": int(args.max_tracks),
        "smooth_sigma_t": float(args.smooth_sigma_t),
        "smooth_sigma_x": float(args.smooth_sigma_x),
        "merge_gap_t": int(args.merge_gap_t),
        "merge_max_dx": float(args.merge_max_dx),
        "min_directionality": float(args.min_directionality),
        "min_net_velocity": float(args.min_net_velocity),
        "merge_min_bridge_prob": float(args.merge_min_bridge_prob),
        "merge_min_bridge_coverage": float(args.merge_min_bridge_coverage),
        "merge_min_embed_sim": float(args.merge_min_embed_sim),
        "keep_path": True,
    }

    sweep_rows = []
    print(f"[info] sweeping {len(thresholds)} threshold(s) over {len(probs)} sample(s)...", flush=True)
    for thr_idx, thr in enumerate(thresholds, start=1):
        t0_thr = time.time()
        counts = {
            "idtp": 0.0,
            "idfp": 0.0,
            "idfn": 0.0,
            "tp_tracks": 0.0,
            "fp_tracks": 0.0,
            "fn_tracks": 0.0,
            "gt_tracks": 0.0,
            "pred_tracks": 0.0,
        }

        for sample_i, (prob, emb_map, gt_tracks) in enumerate(zip(probs, emb_maps, gt_tracks_all), start=1):
            pred_raw = postprocess_prob_to_tracks(
                prob,
                prob_thresh=float(thr),
                embedding_map=emb_map,
                **post_args,
            )
            pred_tracks = _build_pred_tracks(pred_raw)
            sample_counts = _eval_sample_identity(
                gt_tracks,
                pred_tracks,
                tol_px=float(args.tol_px),
                track_iou_thresh=float(args.track_iou_thresh),
                track_min_match_frames=int(args.track_min_match_frames),
            )
            for k in counts:
                counts[k] += float(sample_counts[k])
            if (sample_i % log_every) == 0 or sample_i == len(probs):
                print(
                    f"[info] thr {thr_idx}/{len(thresholds)}={thr:.3f}: "
                    f"{sample_i}/{len(probs)} sample(s)",
                    flush=True,
                )

        sweep_rows.append(
            {
                "threshold": float(thr),
                "summary": _summary_from_counts(counts),
            }
        )
        print(
            f"[info] finished thr {thr_idx}/{len(thresholds)}={thr:.3f} "
            f"in {(time.time() - t0_thr):.1f}s",
            flush=True,
        )

    best_f1 = max(
        sweep_rows,
        key=lambda r: (
            float(r["summary"]["id_f1"]),
            float(r["summary"]["id_recall"]),
            float(r["summary"]["id_precision"]),
        ),
    )
    recall_candidates = [
        r for r in sweep_rows
        if float(r["summary"]["id_precision"]) >= recall_min_precision
    ]
    if not recall_candidates:
        recall_candidates = sweep_rows
    best_recall = max(
        recall_candidates,
        key=lambda r: (
            float(r["summary"]["id_recall"]),
            float(r["summary"]["id_f1"]),
            float(r["summary"]["id_precision"]),
        ),
    )

    payload = {
        "n": len(probs),
        "thresholds": thresholds,
        "rows": sweep_rows,
        "best": {
            "by_id_f1": best_f1,
            "by_id_recall": best_recall,
            "recall_min_precision": recall_min_precision,
        },
        "notes": {
            "skipped_no_meta": int(skipped_no_meta),
            "skipped_missing_identity_paths": int(skipped_missing_paths),
            "tol_px": float(args.tol_px),
            "track_iou_thresh": float(args.track_iou_thresh),
            "track_min_match_frames": int(args.track_min_match_frames),
            "postprocess": post_args,
        },
    }

    print(
        f"Loaded {len(probs)} sample(s) with identity paths. Swept {len(thresholds)} threshold(s). "
        f"Skipped no-meta={skipped_no_meta}, missing-paths={skipped_missing_paths}."
    )
    for row in sweep_rows:
        s = row["summary"]
        print(
            "thr={thr:.3f} idf1={idf1:.4f} idr={idr:.4f} idp={idp:.4f} "
            "trk_f1={tf1:.4f} trk_r={tr:.4f} trk_p={tp:.4f}".format(
                thr=float(row["threshold"]),
                idf1=float(s["id_f1"]),
                idr=float(s["id_recall"]),
                idp=float(s["id_precision"]),
                tf1=float(s["track_f1"]),
                tr=float(s["track_recall"]),
                tp=float(s["track_precision"]),
            )
        )

    bf = best_f1["summary"]
    br = best_recall["summary"]
    print(
        "Best by IDF1: thr={thr:.3f} idf1={idf1:.4f} idr={idr:.4f} idp={idp:.4f}".format(
            thr=float(best_f1["threshold"]),
            idf1=float(bf["id_f1"]),
            idr=float(bf["id_recall"]),
            idp=float(bf["id_precision"]),
        )
    )
    print(
        "Best by ID recall: thr={thr:.3f} idr={idr:.4f} idf1={idf1:.4f} idp={idp:.4f}".format(
            thr=float(best_recall["threshold"]),
            idr=float(br["id_recall"]),
            idf1=float(br["id_f1"]),
            idp=float(br["id_precision"]),
        )
    )
    print(f"[info] total runtime: {(time.time() - t0_all):.1f}s", flush=True)

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


if __name__ == "__main__":
    main()
