from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Iterable

import numpy as np
from scipy import ndimage as ndi

_STRUCTURE_8 = np.ones((3, 3), dtype=bool)


@dataclass
class PostprocessParams:
    prob_thresh: float = 0.35
    min_len: float = 25.0
    vmin: float = 0.15
    simplify_eps: float = 2.0
    max_anchors: int = 12
    min_component_size: int = 18
    min_t_span: int = 16
    min_mean_prob: float = 0.48
    max_tracks: int = 8
    smooth_sigma_t: float = 0.6
    smooth_sigma_x: float = 0.6
    merge_gap_t: int = 24
    merge_max_dx: float = 8.0
    min_directionality: float = 0.10
    min_net_velocity: float = 0.03


def _remove_small(mask: np.ndarray, min_size: int) -> np.ndarray:
    if min_size <= 1:
        return mask
    labels, n = ndi.label(mask, structure=_STRUCTURE_8)
    if n == 0:
        return np.zeros_like(mask, dtype=bool)
    sizes = ndi.sum(mask.astype(np.int32), labels, index=np.arange(1, n + 1))
    keep = np.zeros(n + 1, dtype=bool)
    keep[1:] = sizes >= min_size
    return keep[labels]


def _skeletonize(mask: np.ndarray) -> np.ndarray:
    try:
        from skimage.morphology import skeletonize  # type: ignore

        return skeletonize(mask)
    except Exception:
        return mask


def _sample_embedding_vector(
    emb_map: np.ndarray | None,
    t: int,
    x: float,
    *,
    x_radius: int = 1,
) -> np.ndarray | None:
    if emb_map is None:
        return None
    arr = np.asarray(emb_map, dtype=np.float32)
    if arr.ndim != 3:
        return None
    d, tdim, xdim = arr.shape
    if t < 0 or t >= tdim or d <= 0:
        return None
    xc = int(round(float(x)))
    x0 = max(0, xc - int(max(0, x_radius)))
    x1 = min(xdim, xc + int(max(0, x_radius)) + 1)
    if x0 >= x1:
        return None
    vec = np.mean(arr[:, int(t), x0:x1], axis=1).astype(np.float32, copy=False)
    n = float(np.linalg.norm(vec))
    if n <= 1e-6:
        return None
    return (vec / n).astype(np.float32, copy=False)


def _component_to_path(
    component_mask: np.ndarray,
    prob: np.ndarray | None = None,
    emb_map: np.ndarray | None = None,
) -> list[tuple[float, float]]:
    ys, xs = np.nonzero(component_mask)
    if ys.size == 0:
        return []
    unique_t = np.unique(ys)

    # Coarse slope prior from row-wise weighted centers.
    row_ts: list[float] = []
    row_xs: list[float] = []
    for t in unique_t:
        row_sel = ys == t
        xvals = xs[row_sel]
        if xvals.size == 0:
            continue
        xv = xvals.astype(np.float32)
        wv = np.ones_like(xv, dtype=np.float32)
        if prob is not None:
            row_prob = np.asarray(prob[int(t), xvals], dtype=np.float32)
            if row_prob.size == xvals.size and np.any(row_prob > 0):
                wv = row_prob.astype(np.float32)
        wsum = float(np.sum(wv))
        x_center = float(np.sum(xv * wv) / wsum) if wsum > 1e-9 else float(np.mean(xv))
        row_ts.append(float(t))
        row_xs.append(float(x_center))
    slope_guess = 0.0
    if len(row_ts) >= 2:
        dt_all = float(row_ts[-1] - row_ts[0])
        if abs(dt_all) > 1e-6:
            slope_guess = float((row_xs[-1] - row_xs[0]) / dt_all)
    slope_guess = float(np.clip(slope_guess, -20.0, 20.0))

    rows: list[tuple[int, float]] = []
    prev_t: int | None = None
    prev_x: float | None = None
    prev_v: float = float(slope_guess)
    prev_e: np.ndarray | None = None
    max_fill_gap = 3

    for t in unique_t:
        row_sel = ys == t
        xvals = xs[row_sel]
        if xvals.size == 0:
            continue

        ti = int(t)
        xv = xvals.astype(np.float32)
        wv = np.ones_like(xv, dtype=np.float32)
        if prob is not None:
            row_prob = np.asarray(prob[ti, xvals], dtype=np.float32)
            if row_prob.size == xvals.size and np.any(row_prob > 0):
                wv = row_prob.astype(np.float32)

        if prev_x is None:
            # Seed on dominant local ridge in the first row.
            peak = float(np.max(wv))
            keep = wv >= float(0.88 * peak)
            if np.any(keep):
                xc = xv[keep]
                wc = wv[keep]
                wsum = float(np.sum(wc))
                x_center = float(np.sum(xc * wc) / wsum) if wsum > 1e-9 else float(np.mean(xvals))
            else:
                x_center = float(xvals[int(np.argmax(wv))])
            rows.append((ti, x_center))
            prev_t = ti
            prev_x = x_center
            prev_v = float(slope_guess)
            prev_e = _sample_embedding_vector(emb_map, ti, x_center, x_radius=1)
            continue

        assert prev_t is not None
        dt = max(1, ti - int(prev_t))
        speed_ref = float(max(abs(prev_v), abs(slope_guess)))
        pred_x = float(prev_x + prev_v * float(dt))
        gate = float(np.clip(3.5 + 1.8 * speed_ref * float(dt), 3.5, 28.0))

        dist = np.abs(xv - pred_x)
        keep = dist <= gate
        pred2 = float(prev_x + slope_guess * float(dt))
        dist2 = np.abs(xv - pred2)
        if not np.any(keep) and abs(slope_guess) > 1.5:
            # Reacquire using global slope prior for fast tracks.
            gate2 = float(np.clip(5.0 + 1.8 * abs(slope_guess) * float(dt), 6.0, 32.0))
            keep = dist2 <= gate2
            if np.any(keep):
                pred_x = pred2
                dist = dist2
                gate = gate2
        if not np.any(keep):
            keep = dist <= max(2.2 * gate, 16.0)
        if not np.any(keep):
            # Thin obvious ridges often drift slightly outside the predicted gate.
            # Reacquire with a softer distance penalty before abandoning the row.
            alt_dist = np.minimum(dist, dist2)
            relaxed_gate = float(max(3.0 * gate, 20.0))
            keep = alt_dist <= relaxed_gate
            if np.any(keep):
                dist = alt_dist
                # Blend local velocity and global slope guesses so reacquisition
                # does not overcommit to either one alone.
                pred_x = float((0.65 * pred_x) + (0.35 * pred2))
        if not np.any(keep):
            score_all = (2.4 * wv) - (0.10 * dist) - (0.05 * dist2)
            if prev_e is not None:
                sims = np.zeros((xv.shape[0],), dtype=np.float32)
                for ci in range(xv.shape[0]):
                    ce = _sample_embedding_vector(emb_map, ti, float(xv[ci]), x_radius=1)
                    if ce is not None:
                        sims[ci] = float(np.dot(prev_e, ce))
                score_all = score_all + (0.90 * sims)
            best_all = int(np.argmax(score_all))
            best_dist = float(min(dist[best_all], dist2[best_all]))
            # Still reject distant branch jumps, but allow a strong local ridge
            # to reacquire after a brief mismatch in the predicted position.
            if best_dist > float(max(3.4 * gate, 24.0)):
                continue
            keep = np.zeros_like(xv, dtype=bool)
            keep[best_all] = True
            pred_x = float(xv[best_all])
            dist = np.abs(xv - pred_x)

        xc = xv[keep]
        wc = wv[keep]
        dc = np.abs(xc - pred_x)
        # Continuity-constrained ridge pick: prefer high probability near predicted position.
        score = (2.0 * wc) - (0.25 * dc)
        if prev_e is not None:
            sims = np.zeros((xc.shape[0],), dtype=np.float32)
            for ci in range(xc.shape[0]):
                ce = _sample_embedding_vector(emb_map, ti, float(xc[ci]), x_radius=1)
                if ce is not None:
                    sims[ci] = float(np.dot(prev_e, ce))
            # Embedding continuity helps keep identity through crossings.
            score = score + (0.90 * sims)
        best = int(np.argmax(score))
        x_center = float(xc[best])

        if dt > 1 and dt <= int(max_fill_gap):
            last_t = int(prev_t)
            last_x = float(prev_x)
            for fill_t in range(last_t + 1, ti):
                alpha = float(fill_t - last_t) / float(max(1, ti - last_t))
                fill_x = float((1.0 - alpha) * last_x + alpha * x_center)
                rows.append((int(fill_t), fill_x))

        inst_v = float((x_center - prev_x) / float(dt))
        prev_v = float(np.clip((0.35 * prev_v) + (0.65 * inst_v), -20.0, 20.0))
        curr_e = _sample_embedding_vector(emb_map, ti, x_center, x_radius=1)
        if curr_e is not None:
            if prev_e is None:
                prev_e = curr_e
            else:
                mix = (0.65 * prev_e) + (0.35 * curr_e)
                n = float(np.linalg.norm(mix))
                if n > 1e-6:
                    prev_e = (mix / n).astype(np.float32, copy=False)
        rows.append((ti, x_center))
        prev_t = ti
        prev_x = x_center

    if not rows:
        return []

    # Suppress single-frame zig-zag outliers that deviate from local linear trend.
    if len(rows) >= 3:
        ts = np.asarray([r[0] for r in rows], dtype=np.float32)
        xs_path = np.asarray([r[1] for r in rows], dtype=np.float32)
        for i in range(1, len(rows) - 1):
            t0, t1, t2 = ts[i - 1], ts[i], ts[i + 1]
            if t2 <= t0:
                continue
            a = float((t1 - t0) / max(1e-6, (t2 - t0)))
            x_hat = float((1.0 - a) * xs_path[i - 1] + a * xs_path[i + 1])
            if abs(float(xs_path[i] - x_hat)) > 3.0 and abs(float(xs_path[i + 1] - xs_path[i - 1])) < 4.0:
                xs_path[i] = x_hat
        rows = [(int(ts[i]), float(xs_path[i])) for i in range(len(rows))]

    pts = [(float(t), float(x)) for t, x in rows]
    pts.sort(key=lambda p: p[0])
    return pts


def _skeleton_component_to_path(
    component_mask: np.ndarray,
    prob: np.ndarray | None = None,
) -> list[tuple[float, float]]:
    ys, xs = np.nonzero(component_mask)
    if ys.size == 0:
        return []

    unique_t = np.unique(ys)
    rows: list[tuple[int, float]] = []
    prev_t: int | None = None
    prev_x: float | None = None
    max_fill_gap = 4

    for t in unique_t:
        ti = int(t)
        row_x = xs[ys == t]
        if row_x.size == 0:
            continue

        xvals = row_x.astype(np.float32)
        if prob is not None:
            row_prob = np.asarray(prob[ti, row_x], dtype=np.float32)
            if row_prob.size == row_x.size and np.any(row_prob > 0):
                peak = float(np.max(row_prob))
                keep = row_prob >= float(max(0.85 * peak, peak - 0.03))
                if np.any(keep):
                    sel_x = xvals[keep]
                    sel_w = row_prob[keep].astype(np.float32)
                    wsum = float(np.sum(sel_w))
                    x_center = float(np.sum(sel_x * sel_w) / wsum) if wsum > 1e-9 else float(np.mean(sel_x))
                else:
                    x_center = float(xvals[int(np.argmax(row_prob))])
            else:
                x_center = float(np.median(xvals))
        else:
            x_center = float(np.median(xvals))

        if prev_t is not None and prev_x is not None:
            dt = int(ti - prev_t)
            if dt > 1 and dt <= int(max_fill_gap):
                for fill_t in range(prev_t + 1, ti):
                    alpha = float(fill_t - prev_t) / float(max(1, ti - prev_t))
                    fill_x = float((1.0 - alpha) * prev_x + alpha * x_center)
                    rows.append((int(fill_t), fill_x))

        rows.append((ti, x_center))
        prev_t = ti
        prev_x = x_center

    if len(rows) >= 3:
        ts = np.asarray([r[0] for r in rows], dtype=np.float32)
        xs_path = np.asarray([r[1] for r in rows], dtype=np.float32)
        for i in range(1, len(rows) - 1):
            t0, t1, t2 = ts[i - 1], ts[i], ts[i + 1]
            if t2 <= t0:
                continue
            a = float((t1 - t0) / max(1e-6, (t2 - t0)))
            x_hat = float((1.0 - a) * xs_path[i - 1] + a * xs_path[i + 1])
            if abs(float(xs_path[i] - x_hat)) > 3.0 and abs(float(xs_path[i + 1] - xs_path[i - 1])) < 5.0:
                xs_path[i] = x_hat
        rows = [(int(ts[i]), float(xs_path[i])) for i in range(len(rows))]

    pts = [(float(t), float(x)) for t, x in rows]
    pts.sort(key=lambda p: p[0])
    return pts


def _skeleton_label_dp_path(
    component_mask: np.ndarray,
    prob: np.ndarray | None = None,
) -> list[tuple[float, float]]:
    ys, xs = np.nonzero(component_mask)
    if ys.size == 0:
        return []

    unique_t = np.unique(ys)
    if unique_t.size < 2:
        return []

    row_states: list[tuple[int, np.ndarray, np.ndarray]] = []
    row_centers: list[float] = []
    for t in unique_t:
        ti = int(t)
        row_x = np.unique(xs[ys == t].astype(np.int32))
        if row_x.size == 0:
            continue
        row_w = np.ones((row_x.size,), dtype=np.float32)
        if prob is not None:
            row_prob = np.asarray(prob[ti, row_x], dtype=np.float32)
            if row_prob.size == row_x.size:
                row_w = np.maximum(row_prob.astype(np.float32), 1e-4)
        row_states.append((ti, row_x, row_w))
        wsum = float(np.sum(row_w))
        row_centers.append(float(np.sum(row_x.astype(np.float32) * row_w) / wsum) if wsum > 1e-9 else float(np.mean(row_x)))

    if len(row_states) < 2:
        return []

    slope_guess = 0.0
    dt_all = float(row_states[-1][0] - row_states[0][0])
    if abs(dt_all) > 1e-6:
        slope_guess = float((row_centers[-1] - row_centers[0]) / dt_all)
    slope_guess = float(np.clip(slope_guess, -20.0, 20.0))

    scores: list[np.ndarray] = []
    parents: list[np.ndarray] = []

    t0, xs0, w0 = row_states[0]
    score0 = (3.0 * w0.astype(np.float32)) - (0.02 * np.abs(xs0.astype(np.float32) - float(row_centers[0])))
    scores.append(score0.astype(np.float32))
    parents.append(np.full((xs0.shape[0],), -1, dtype=np.int32))

    for row_idx in range(1, len(row_states)):
        ti, xcur, wcur = row_states[row_idx]
        tp, xprev, _wprev = row_states[row_idx - 1]
        dt = max(1, int(ti - tp))
        xcur_f = xcur.astype(np.float32)
        xprev_f = xprev.astype(np.float32)

        pred = xprev_f[:, None] + float(slope_guess) * float(dt)
        dx = np.abs(xcur_f[None, :] - xprev_f[:, None])
        dx_pred = np.abs(xcur_f[None, :] - pred)
        trans = (0.18 * dx) + (0.28 * dx_pred)
        if dt > 1:
            trans = trans + float(0.10 * (dt - 1))

        cand = scores[-1][:, None] - trans
        best_prev = np.argmax(cand, axis=0).astype(np.int32)
        best_score = cand[best_prev, np.arange(xcur.shape[0])]
        local = (3.0 * wcur.astype(np.float32)) - (0.015 * np.abs(xcur_f - float(row_centers[row_idx])))
        scores.append((best_score + local).astype(np.float32))
        parents.append(best_prev)

    end_idx = int(np.argmax(scores[-1]))
    chosen: list[tuple[int, float]] = []
    for row_idx in range(len(row_states) - 1, -1, -1):
        ti, xrow, _wrow = row_states[row_idx]
        end_idx = int(np.clip(end_idx, 0, max(0, xrow.shape[0] - 1)))
        chosen.append((int(ti), float(xrow[end_idx])))
        prev_idx = int(parents[row_idx][end_idx]) if row_idx > 0 else -1
        end_idx = prev_idx
    chosen.reverse()

    rows: list[tuple[int, float]] = []
    for i, (ti, xi) in enumerate(chosen):
        if rows:
            last_t, last_x = rows[-1]
            gap = int(ti - last_t)
            if gap > 1 and gap <= 4:
                for fill_t in range(last_t + 1, ti):
                    alpha = float(fill_t - last_t) / float(max(1, ti - last_t))
                    fill_x = float((1.0 - alpha) * last_x + alpha * xi)
                    rows.append((int(fill_t), fill_x))
        rows.append((int(ti), float(xi)))

    if len(rows) >= 3:
        ts = np.asarray([r[0] for r in rows], dtype=np.float32)
        xs_path = np.asarray([r[1] for r in rows], dtype=np.float32)
        for i in range(1, len(rows) - 1):
            t_prev, t_cur, t_next = ts[i - 1], ts[i], ts[i + 1]
            if t_next <= t_prev:
                continue
            a = float((t_cur - t_prev) / max(1e-6, (t_next - t_prev)))
            x_hat = float((1.0 - a) * xs_path[i - 1] + a * xs_path[i + 1])
            if abs(float(xs_path[i] - x_hat)) > 2.5 and abs(float(xs_path[i + 1] - xs_path[i - 1])) < 5.0:
                xs_path[i] = x_hat
        rows = [(int(ts[i]), float(xs_path[i])) for i in range(len(rows))]

    pts = [(float(t), float(x)) for t, x in rows]
    pts.sort(key=lambda p: p[0])
    return pts


def _skeleton_branch_candidate_masks(component_mask: np.ndarray) -> list[np.ndarray]:
    if not np.any(component_mask):
        return []
    mask_i = component_mask.astype(np.int32, copy=False)
    neighbor_count = ndi.convolve(mask_i, _STRUCTURE_8.astype(np.int32), mode="constant", cval=0) - mask_i
    branch_mask = np.logical_and(component_mask, neighbor_count >= 3)
    if not np.any(branch_mask):
        return []

    pruned = np.logical_and(component_mask, ~branch_mask)
    labels, n = ndi.label(pruned, structure=_STRUCTURE_8)
    if n <= 0:
        return []

    pieces: list[np.ndarray] = []
    for comp_id in range(1, n + 1):
        piece = labels == comp_id
        if not np.any(piece):
            continue
        rows = np.nonzero(np.any(piece, axis=1))[0]
        if rows.size < 2:
            continue
        expanded = np.logical_and(ndi.binary_dilation(piece, structure=_STRUCTURE_8, iterations=1), component_mask)
        exp_rows = np.nonzero(np.any(expanded, axis=1))[0]
        if exp_rows.size < 2:
            continue
        duplicate = False
        for prev in pieces:
            inter = float(np.sum(np.logical_and(prev, expanded)))
            union = float(np.sum(np.logical_or(prev, expanded)))
            if union > 0.0 and inter / union >= 0.85:
                duplicate = True
                break
        if not duplicate:
            pieces.append(expanded)
    return pieces


def _component_row_span_stats(component_mask: np.ndarray) -> tuple[float, float]:
    ys, xs = np.nonzero(component_mask)
    if ys.size == 0:
        return 0.0, 0.0
    unique_t = np.unique(ys)
    spans: list[float] = []
    for t in unique_t:
        row = xs[ys == t]
        if row.size == 0:
            continue
        spans.append(float(np.max(row) - np.min(row)))
    if not spans:
        return 0.0, 0.0
    arr = np.asarray(spans, dtype=np.float32)
    return float(np.median(arr)), float(np.percentile(arr, 90.0))


def _adaptive_threshold(prob: np.ndarray, base_thresh: float) -> float:
    # Keep probability threshold literal and predictable.
    return float(np.clip(base_thresh, 0.0, 0.9999))


def _build_mask(
    prob: np.ndarray,
    prob_thresh: float,
    min_component_size: int,
    debug: dict | None = None,
) -> np.ndarray:
    thresh = _adaptive_threshold(prob, float(prob_thresh))
    high_mask = prob >= thresh

    # Rescue thin ridge peaks that sit just below the global threshold without
    # admitting broad low-confidence noise everywhere. These seeds help
    # obviously track-like local maxima survive into the component stage.
    ridge_seed_thresh = float(max(0.72, thresh - 0.05))
    row_local_max = prob >= ndi.maximum_filter(prob, size=(1, 3))
    ridge_seeds = np.logical_and(row_local_max, prob >= ridge_seed_thresh)
    seed_mask = np.logical_or(high_mask, ridge_seeds)

    # Use a seeded hysteresis mask rather than a single hard cutoff. This keeps
    # a strong ridge connected through nearby slightly weaker pixels without
    # globally lowering the threshold and flooding the image with noise.
    low_thresh = float(max(0.64, thresh - 0.09))
    if low_thresh < thresh and np.any(seed_mask):
        low_mask = prob >= low_thresh
        labels, n = ndi.label(low_mask, structure=_STRUCTURE_8)
        if n > 0:
            seed_labels = np.unique(labels[seed_mask])
            seed_labels = seed_labels[seed_labels > 0]
            keep = np.zeros(n + 1, dtype=bool)
            keep[seed_labels] = True
            mask = keep[labels]
        else:
            mask = seed_mask.copy()
    else:
        mask = seed_mask.copy()

    # Keep topology stable across nearby thresholds instead of switching to a
    # completely different ridge-extraction regime once density crosses a cutoff.
    mask = ndi.binary_closing(mask, structure=np.ones((3, 3), dtype=bool))
    if float(np.mean(mask)) > 0.65:
        mask = ndi.binary_opening(
            mask,
            structure=np.array([[0, 1, 0], [1, 1, 1], [0, 1, 0]], dtype=bool),
        )

    mask = _remove_small(mask, min_size=int(min_component_size))
    if isinstance(debug, dict):
        debug["mask_stats"] = {
            "high_thresh": float(thresh),
            "seed_thresh": float(ridge_seed_thresh),
            "low_thresh": float(low_thresh),
            "high_pixels": int(np.sum(high_mask)),
            "seed_pixels": int(np.sum(ridge_seeds)),
            "mask_pixels": int(np.sum(mask)),
            "mask_density": float(np.mean(mask)),
        }
    return mask


def _polyline_len(points: Iterable[tuple[float, float]]) -> float:
    pts = np.asarray(list(points), dtype=np.float32)
    if pts.shape[0] < 2:
        return 0.0
    dif = np.diff(pts, axis=0)
    seg = np.sqrt(np.sum(dif**2, axis=1))
    return float(np.sum(seg))


def _mean_abs_velocity(points: Iterable[tuple[float, float]]) -> float:
    pts = np.asarray(list(points), dtype=np.float32)
    if pts.shape[0] < 2:
        return 0.0
    dt = np.diff(pts[:, 0])
    dx = np.diff(pts[:, 1])
    valid = np.abs(dt) > 1e-6
    if not np.any(valid):
        return 0.0
    return float(np.mean(np.abs(dx[valid] / dt[valid])))


def _velocity_quantile(points: Iterable[tuple[float, float]], q: float = 75.0) -> float:
    pts = np.asarray(list(points), dtype=np.float32)
    if pts.shape[0] < 2:
        return 0.0
    dt = np.diff(pts[:, 0])
    dx = np.diff(pts[:, 1])
    valid = np.abs(dt) > 1e-6
    if not np.any(valid):
        return 0.0
    v = np.abs(dx[valid] / dt[valid])
    return float(np.percentile(v, float(q)))


def _net_velocity(points: Iterable[tuple[float, float]]) -> float:
    pts = np.asarray(list(points), dtype=np.float32)
    if pts.shape[0] < 2:
        return 0.0
    dt = float(pts[-1, 0] - pts[0, 0])
    if abs(dt) <= 1e-6:
        return 0.0
    dx = float(pts[-1, 1] - pts[0, 1])
    return float(abs(dx / dt))


def _directionality(points: Iterable[tuple[float, float]]) -> float:
    pts = np.asarray(list(points), dtype=np.float32)
    if pts.shape[0] < 2:
        return 0.0
    dx = np.diff(pts[:, 1])
    net = abs(float(np.sum(dx)))
    total = float(np.sum(np.abs(dx)))
    if total <= 1e-6:
        return 0.0
    return float(net / total)


def _motion_strength(points: Iterable[tuple[float, float]]) -> float:
    # Robust motion proxy for paused tracks: combines average, tail quantile, and net drift.
    return float(max(_mean_abs_velocity(points), _velocity_quantile(points, q=75.0), _net_velocity(points)))


def _endpoint_velocity(points: Iterable[tuple[float, float]], tail_points: int = 6, from_start: bool = False) -> float:
    pts = np.asarray(list(points), dtype=np.float32)
    if pts.shape[0] < 2:
        return 0.0
    n = min(int(max(2, tail_points)), pts.shape[0])
    seg = pts[:n] if from_start else pts[-n:]
    dt = np.diff(seg[:, 0])
    dx = np.diff(seg[:, 1])
    valid = np.abs(dt) > 1e-6
    if not np.any(valid):
        return 0.0
    return float(np.median(dx[valid] / dt[valid]))


def _sign(v: float, tol: float = 0.03) -> int:
    if v > tol:
        return 1
    if v < -tol:
        return -1
    return 0


def _sample_prob_local_max(prob: np.ndarray, t: int, x: float, x_radius: int = 1) -> float:
    if t < 0 or t >= prob.shape[0]:
        return 0.0
    xc = int(round(float(x)))
    x0 = max(0, xc - int(max(0, x_radius)))
    x1 = min(prob.shape[1], xc + int(max(0, x_radius)) + 1)
    if x0 >= x1:
        return 0.0
    return float(np.max(prob[t, x0:x1]))


def _bridge_prob_values(
    prob: np.ndarray,
    t0: float,
    x0: float,
    t1: float,
    x1: float,
    *,
    x_radius: int = 1,
) -> np.ndarray:
    t0i = int(round(float(t0)))
    t1i = int(round(float(t1)))
    if t1i <= t0i + 1:
        return np.zeros((0,), dtype=np.float32)
    dt = float(t1i - t0i)
    vals = []
    for t in range(t0i + 1, t1i):
        alpha = float(t - t0i) / dt
        x = float((1.0 - alpha) * x0 + alpha * x1)
        vals.append(_sample_prob_local_max(prob, t, x, x_radius=x_radius))
    if not vals:
        return np.zeros((0,), dtype=np.float32)
    return np.asarray(vals, dtype=np.float32)


def _point_line_distance(p: np.ndarray, a: np.ndarray, b: np.ndarray) -> float:
    ab = b - a
    denom = float(np.dot(ab, ab))
    if denom <= 1e-12:
        return float(np.linalg.norm(p - a))
    t = float(np.dot(p - a, ab) / denom)
    t = min(1.0, max(0.0, t))
    proj = a + t * ab
    return float(np.linalg.norm(p - proj))


def _rdp(points: list[tuple[float, float]], epsilon: float) -> list[tuple[float, float]]:
    if len(points) < 3 or epsilon <= 0:
        return points

    arr = np.asarray(points, dtype=np.float32)
    a = arr[0]
    b = arr[-1]
    dmax = -1.0
    idx = -1
    for i in range(1, len(points) - 1):
        d = _point_line_distance(arr[i], a, b)
        if d > dmax:
            dmax = d
            idx = i

    if dmax > epsilon and idx > 0:
        left = _rdp(points[: idx + 1], epsilon)
        right = _rdp(points[idx:], epsilon)
        return left[:-1] + right
    return [points[0], points[-1]]


def _cap_anchor_count(points: list[tuple[float, float]], max_anchors: int) -> list[tuple[float, float]]:
    if max_anchors <= 0 or len(points) <= max_anchors:
        return points
    idx = np.linspace(0, len(points) - 1, num=max_anchors, endpoint=True)
    idx = np.unique(np.round(idx).astype(int))
    reduced = [points[i] for i in idx]
    if reduced[0] != points[0]:
        reduced = [points[0]] + reduced
    if reduced[-1] != points[-1]:
        reduced = reduced + [points[-1]]
    return reduced


def _track_path(track: dict) -> np.ndarray:
    raw = track.get("path")
    if raw is not None:
        arr = np.asarray(raw, dtype=np.float32)
        if arr.ndim == 2 and arr.shape[1] == 2 and arr.shape[0] >= 2:
            return arr
    anchors = track.get("anchors", []) or []
    if len(anchors) < 2:
        return np.zeros((0, 2), dtype=np.float32)
    pts = [(float(a.get("t", 0)), float(a.get("x", 0.0))) for a in anchors]
    return np.asarray(pts, dtype=np.float32)


def _clone_track_debug(track: dict) -> dict:
    out: dict = {}
    for key, value in track.items():
        if key == "anchors" and isinstance(value, list):
            out[key] = [{"t": int(round(float(a.get("t", 0)))), "x": float(a.get("x", 0.0))} for a in value if isinstance(a, dict)]
        elif key == "path":
            arr = np.asarray(value, dtype=np.float32)
            if arr.ndim == 2 and arr.shape[1] == 2:
                out[key] = [(float(t), float(x)) for t, x in arr]
        elif isinstance(value, np.generic):
            out[key] = value.item()
        else:
            out[key] = value
    return out


def _bump_debug_count(debug_counts: dict | None, key: str) -> None:
    if not isinstance(debug_counts, dict):
        return
    debug_counts[key] = int(debug_counts.get(key, 0)) + 1


_REJECT_REASON_CODES = {
    "accepted": 1,
    "reject_row_span_median": 2,
    "reject_row_span_p90": 3,
    "reject_path_short": 4,
    "reject_t_span": 5,
    "reject_length": 6,
    "reject_vmin": 7,
    "reject_directionality": 8,
    "reject_net_velocity": 9,
    "reject_mean_prob": 10,
    "reject_simplified_short": 11,
}


def _mark_reason_map(reason_map: np.ndarray | None, component_mask: np.ndarray, reason_key: str) -> None:
    if reason_map is None:
        return
    code = int(_REJECT_REASON_CODES.get(reason_key, 0))
    if code <= 0:
        return
    reason_map[component_mask] = code


def _collapse_path_unique_t(path: np.ndarray) -> np.ndarray:
    if path.size == 0:
        return path.reshape(0, 2)
    arr = np.asarray(path, dtype=np.float32)
    order = np.argsort(arr[:, 0])
    arr = arr[order]
    t_int = np.round(arr[:, 0]).astype(np.int32)
    uniq = np.unique(t_int)
    out = np.zeros((len(uniq), 2), dtype=np.float32)
    for i, t in enumerate(uniq):
        xs = arr[t_int == t, 1]
        out[i, 0] = float(t)
        out[i, 1] = float(np.mean(xs))
    return out


def _build_track_from_path(
    path_arr: np.ndarray,
    sources: list[dict],
    *,
    simplify_eps: float,
    max_anchors: int,
) -> dict | None:
    path_arr = _collapse_path_unique_t(path_arr)
    if path_arr.shape[0] < 2:
        return None
    path = [(float(t), float(x)) for t, x in path_arr]
    t_span = float(path[-1][0] - path[0][0])
    length = _polyline_len(path)
    mean_abs_dxdt = _mean_abs_velocity(path)
    motion = _motion_strength(path)
    net_vel = _net_velocity(path)
    dir_score = _directionality(path)
    weighted = 0.0
    wsum = 0.0
    for tr in sources:
        w = float(max(1.0, tr.get("t_span", 1.0)))
        weighted += w * float(tr.get("mean_prob", 0.0))
        wsum += w
    mean_prob = float(weighted / max(1e-6, wsum))
    simplified = _rdp(path, epsilon=float(simplify_eps))
    simplified = _cap_anchor_count(simplified, int(max_anchors))
    if len(simplified) < 2:
        return None
    anchors = [{"t": int(round(t)), "x": float(x)} for t, x in simplified]
    anchors.sort(key=lambda a: a["t"])
    score = float(mean_prob * (length**0.72) * (1.0 + 0.18 * motion) * (1.0 + 0.008 * t_span))
    return {
        "anchors": anchors,
        "path": path,
        "length_px": float(length),
        "mean_abs_dxdt": float(mean_abs_dxdt),
        "motion_strength": float(motion),
        "net_velocity": float(net_vel),
        "directionality": float(dir_score),
        "mean_prob": float(mean_prob),
        "t_span": float(t_span),
        "score": score,
    }


def _try_merge_pair(
    a: dict,
    b: dict,
    *,
    max_gap_t: int,
    max_link_dx: float,
    prob: np.ndarray | None = None,
    min_bridge_prob: float = 0.0,
    min_bridge_coverage: float = 0.0,
    emb_map: np.ndarray | None = None,
    min_embed_sim: float | None = None,
) -> np.ndarray | None:
    pa = _track_path(a)
    pb = _track_path(b)
    if pa.shape[0] < 2 or pb.shape[0] < 2:
        return None

    pa = _collapse_path_unique_t(pa)
    pb = _collapse_path_unique_t(pb)
    a_end_t, a_end_x = float(pa[-1, 0]), float(pa[-1, 1])
    b_start_t, b_start_x = float(pb[0, 0]), float(pb[0, 1])
    gap = int(round(b_start_t - a_end_t))
    if gap < 0 or gap > int(max_gap_t):
        return None

    va = _endpoint_velocity(pa, tail_points=6, from_start=False)
    vb = _endpoint_velocity(pb, tail_points=6, from_start=True)
    sa = _sign(va)
    sb = _sign(vb)
    if sa != 0 and sb != 0 and sa != sb:
        return None
    # Reject merges with strongly incompatible endpoint velocities.
    if abs(float(va) - float(vb)) > (0.35 + 0.06 * max(0, gap)):
        return None

    pred_x = float(a_end_x + va * max(0, gap))
    dx_raw = abs(b_start_x - a_end_x)
    dx_pred = abs(b_start_x - pred_x)
    tol = float(max_link_dx + 0.18 * max(0, gap))
    if min(dx_raw, dx_pred) > tol:
        return None

    if emb_map is not None and min_embed_sim is not None:
        ea = _sample_embedding_vector(emb_map, int(round(a_end_t)), a_end_x, x_radius=1)
        eb = _sample_embedding_vector(emb_map, int(round(b_start_t)), b_start_x, x_radius=1)
        if ea is not None and eb is not None:
            sim = float(np.dot(ea, eb))
            if sim < float(min_embed_sim):
                return None

    # For non-adjacent segments, require probability support between endpoints.
    if prob is not None and gap > 1:
        bridge_vals = _bridge_prob_values(prob, a_end_t, a_end_x, b_start_t, b_start_x, x_radius=1)
        if bridge_vals.size > 0:
            floor = float(max(0.0, min_bridge_prob))
            coverage = float(np.mean(bridge_vals >= floor)) if floor > 0 else 1.0
            mean_v = float(np.mean(bridge_vals))
            peak_v = float(np.max(bridge_vals))
            need_cov = float(max(0.0, min(1.0, min_bridge_coverage)))
            if floor > 0.0:
                weak_bridge = (
                    coverage < need_cov
                    and mean_v < (0.9 * floor)
                    and peak_v < (1.15 * floor)
                )
                if weak_bridge:
                    return None

    merged = pa
    if gap > 1:
        t0 = int(round(a_end_t))
        t1 = int(round(b_start_t))
        bridge_t = np.arange(t0 + 1, t1, dtype=np.int32)
        if bridge_t.size > 0:
            bridge_x = np.linspace(a_end_x, b_start_x, num=bridge_t.size + 2, endpoint=True)[1:-1]
            bridge = np.column_stack([bridge_t.astype(np.float32), bridge_x.astype(np.float32)])
            merged = np.vstack([merged, bridge])

    append = pb[pb[:, 0] > merged[-1, 0]]
    if append.size > 0:
        merged = np.vstack([merged, append])

    merged = _collapse_path_unique_t(merged)
    return merged


def _merge_split_tracks(
    tracks: list[dict],
    *,
    simplify_eps: float,
    max_anchors: int,
    max_gap_t: int,
    max_link_dx: float,
    prob: np.ndarray | None = None,
    min_bridge_prob: float = 0.0,
    min_bridge_coverage: float = 0.0,
    emb_map: np.ndarray | None = None,
    min_embed_sim: float | None = None,
) -> list[dict]:
    if len(tracks) < 2:
        return tracks

    # Runtime guard: for very dense candidate sets, use the previous
    # linear-greedy strategy to avoid cubic blow-ups.
    if len(tracks) > 72:
        ordered = sorted(tracks, key=lambda tr: float(_track_path(tr)[0, 0]) if _track_path(tr).shape[0] else 1e9)
        used = [False] * len(ordered)
        merged_tracks: list[dict] = []

        for i, tr in enumerate(ordered):
            if used[i]:
                continue
            used[i] = True
            current = tr

            while True:
                best_j = -1
                best_cost = float("inf")
                best_path = None
                curr_path = _track_path(current)
                if curr_path.shape[0] < 2:
                    break

                for j, cand in enumerate(ordered):
                    if used[j]:
                        continue
                    merged_path = _try_merge_pair(
                        current,
                        cand,
                        max_gap_t=int(max_gap_t),
                        max_link_dx=float(max_link_dx),
                        prob=prob,
                        min_bridge_prob=float(min_bridge_prob),
                        min_bridge_coverage=float(min_bridge_coverage),
                        emb_map=emb_map,
                        min_embed_sim=min_embed_sim,
                    )
                    if merged_path is None:
                        continue
                    cand_path = _track_path(cand)
                    if cand_path.shape[0] < 1:
                        continue
                    gap = max(0.0, float(cand_path[0, 0] - curr_path[-1, 0]))
                    dx = abs(float(cand_path[0, 1] - curr_path[-1, 1]))
                    cost = float(dx + 0.2 * gap)
                    if cost < best_cost:
                        best_cost = cost
                        best_j = j
                        best_path = merged_path

                if best_j < 0 or best_path is None:
                    break

                merged = _build_track_from_path(
                    best_path,
                    [current, ordered[best_j]],
                    simplify_eps=float(simplify_eps),
                    max_anchors=int(max_anchors),
                )
                if merged is None:
                    break
                used[best_j] = True
                current = merged

            merged_tracks.append(current)

        return merged_tracks

    working = list(tracks)

    def _merge_pair_cost(a: dict, b: dict, merged_path: np.ndarray) -> float:
        pa = _collapse_path_unique_t(_track_path(a))
        pb = _collapse_path_unique_t(_track_path(b))
        if pa.shape[0] < 2 or pb.shape[0] < 2:
            return float("inf")

        a_end_t, a_end_x = float(pa[-1, 0]), float(pa[-1, 1])
        b_start_t, b_start_x = float(pb[0, 0]), float(pb[0, 1])
        gap = max(0.0, float(b_start_t - a_end_t))
        va = _endpoint_velocity(pa, tail_points=6, from_start=False)
        vb = _endpoint_velocity(pb, tail_points=6, from_start=True)
        pred_x = float(a_end_x + va * gap)
        dx = min(abs(b_start_x - a_end_x), abs(b_start_x - pred_x))

        cost = float((0.55 * dx) + (0.55 * gap) + (0.45 * abs(va - vb)))

        if prob is not None and gap > 1:
            bridge_vals = _bridge_prob_values(prob, a_end_t, a_end_x, b_start_t, b_start_x, x_radius=1)
            if bridge_vals.size > 0:
                mean_v = float(np.mean(bridge_vals))
                cov = float(np.mean(bridge_vals >= max(0.0, float(min_bridge_prob))))
                cost -= float(2.2 * mean_v + 1.1 * cov)

        if emb_map is not None:
            ea = _sample_embedding_vector(emb_map, int(round(a_end_t)), a_end_x, x_radius=1)
            eb = _sample_embedding_vector(emb_map, int(round(b_start_t)), b_start_x, x_radius=1)
            if ea is not None and eb is not None:
                sim = float(np.dot(ea, eb))
                # Prefer identity-consistent merges strongly.
                cost -= float(1.4 * max(-0.05, sim))

        # Prefer merges that improve total span continuity.
        merged_path = _collapse_path_unique_t(merged_path)
        if merged_path.shape[0] >= 2:
            merged_span = float(merged_path[-1, 0] - merged_path[0, 0])
            base_span = float(a.get("t_span", 0.0)) + float(b.get("t_span", 0.0))
            gain = max(0.0, merged_span - max(float(a.get("t_span", 0.0)), float(b.get("t_span", 0.0))))
            cost -= float(0.015 * gain)
            cost += float(0.003 * max(0.0, base_span - merged_span))

        return float(cost)

    # Global best-first agglomerative merge. This is more stable than
    # per-track greedy chaining when many fragments compete at crossovers.
    while True:
        if len(working) < 2:
            break

        best_i = -1
        best_j = -1
        best_cost = float("inf")
        best_path = None

        for i, a in enumerate(working):
            pa = _collapse_path_unique_t(_track_path(a))
            if pa.shape[0] < 2:
                continue
            a_end_t = float(pa[-1, 0])
            for j, b in enumerate(working):
                if i == j:
                    continue
                pb = _collapse_path_unique_t(_track_path(b))
                if pb.shape[0] < 2:
                    continue
                gap = float(pb[0, 0] - a_end_t)
                if gap < 0.0 or gap > float(max_gap_t):
                    continue
                merged_path = _try_merge_pair(
                    a,
                    b,
                    max_gap_t=int(max_gap_t),
                    max_link_dx=float(max_link_dx),
                    prob=prob,
                    min_bridge_prob=float(min_bridge_prob),
                    min_bridge_coverage=float(min_bridge_coverage),
                    emb_map=emb_map,
                    min_embed_sim=min_embed_sim,
                )
                if merged_path is None:
                    continue

                cost = _merge_pair_cost(a, b, merged_path)
                if cost < best_cost:
                    best_cost = cost
                    best_i = i
                    best_j = j
                    best_path = merged_path

        if best_i < 0 or best_j < 0 or best_path is None:
            break

        a = working[best_i]
        b = working[best_j]
        merged = _build_track_from_path(
            best_path,
            [a, b],
            simplify_eps=float(simplify_eps),
            max_anchors=int(max_anchors),
        )
        if merged is None:
            break

        # Remove higher index first to keep indices stable.
        for idx in sorted([best_i, best_j], reverse=True):
            del working[idx]
        working.append(merged)

    return working


def _extend_single_track_path(
    path_arr: np.ndarray,
    prob: np.ndarray,
    *,
    emb_map: np.ndarray | None = None,
    max_extend_t: int = 18,
    search_dx: int = 6,
    min_prob: float = 0.05,
    max_miss: int = 2,
) -> np.ndarray:
    path = _collapse_path_unique_t(path_arr)
    if path.shape[0] < 2:
        return path

    tdim, xdim = prob.shape

    def _pick_x(
        t: int,
        pred_x: float,
        prev_x: float,
        prev_v: float,
        ref_e: np.ndarray | None,
        *,
        backward: bool = False,
    ) -> tuple[float, np.ndarray | None] | None:
        if t < 0 or t >= tdim:
            return None
        xr = int(max(1, search_dx))
        xc = int(round(float(pred_x)))
        x0 = max(0, xc - xr)
        x1 = min(xdim - 1, xc + xr)
        if x0 > x1:
            return None
        xs = np.arange(x0, x1 + 1, dtype=np.int32)
        pv = prob[int(t), xs].astype(np.float32, copy=False)
        floor = float(max(0.0, min_prob))
        keep = pv >= floor
        if not np.any(keep):
            relaxed = max(0.0, 0.75 * floor)
            keep = pv >= relaxed
            if not np.any(keep):
                return None
        xsf = xs[keep].astype(np.float32)
        psel = pv[keep].astype(np.float32)
        dx_pred = np.abs(xsf - float(pred_x))
        if backward:
            v_err = np.abs((xsf - float(prev_x)) + float(prev_v))
        else:
            v_err = np.abs((xsf - float(prev_x)) - float(prev_v))
        score = (2.5 * psel) - (0.17 * dx_pred) - (0.09 * v_err)

        best_e: np.ndarray | None = None
        if emb_map is not None and ref_e is not None:
            sims = np.zeros((xsf.shape[0],), dtype=np.float32)
            for k in range(xsf.shape[0]):
                ce = _sample_embedding_vector(emb_map, int(t), float(xsf[k]), x_radius=1)
                if ce is not None:
                    sims[k] = float(np.dot(ref_e, ce))
            score = score + (0.75 * sims)

        kbest = int(np.argmax(score))
        x_best = float(xsf[kbest])
        if emb_map is not None:
            best_e = _sample_embedding_vector(emb_map, int(t), x_best, x_radius=1)
        return x_best, best_e

    # Forward extension.
    pts_f = [(float(t), float(x)) for t, x in path]
    vf = _endpoint_velocity(path, tail_points=6, from_start=False)
    tf = int(round(float(path[-1, 0])))
    xf = float(path[-1, 1])
    ef = _sample_embedding_vector(emb_map, tf, xf, x_radius=1)
    miss = 0
    steps = 0
    while tf < (tdim - 1) and steps < int(max_extend_t):
        t_next = tf + 1
        pred_x = float(xf + vf)
        pick = _pick_x(t_next, pred_x, xf, vf, ef, backward=False)
        if pick is None:
            miss += 1
            tf = t_next
            if miss > int(max_miss):
                break
            continue
        x_new, e_new = pick
        dt = max(1, int(t_next - int(round(pts_f[-1][0]))))
        if dt > 1:
            t0, x0 = pts_f[-1]
            for tt in range(int(round(t0)) + 1, int(t_next)):
                a = float(tt - t0) / float(t_next - t0)
                xx = float((1.0 - a) * x0 + a * x_new)
                pts_f.append((float(tt), xx))
        pts_f.append((float(t_next), float(x_new)))
        inst_v = float((x_new - xf) / max(1.0, float(dt)))
        vf = float(np.clip((0.60 * vf) + (0.40 * inst_v), -20.0, 20.0))
        xf = float(x_new)
        tf = int(t_next)
        miss = 0
        steps += 1
        if e_new is not None:
            if ef is None:
                ef = e_new
            else:
                mix = (0.65 * ef) + (0.35 * e_new)
                n = float(np.linalg.norm(mix))
                if n > 1e-6:
                    ef = (mix / n).astype(np.float32, copy=False)

    # Backward extension.
    pts_b = [(float(t), float(x)) for t, x in pts_f]
    pb = np.asarray(pts_b, dtype=np.float32)
    pb = _collapse_path_unique_t(pb)
    vb = _endpoint_velocity(pb, tail_points=6, from_start=True)
    tb = int(round(float(pb[0, 0])))
    xb = float(pb[0, 1])
    eb = _sample_embedding_vector(emb_map, tb, xb, x_radius=1)
    miss = 0
    steps = 0
    prepend: list[tuple[float, float]] = []
    while tb > 0 and steps < int(max_extend_t):
        t_prev = tb - 1
        pred_x = float(xb - vb)
        pick = _pick_x(t_prev, pred_x, xb, vb, eb, backward=True)
        if pick is None:
            miss += 1
            tb = t_prev
            if miss > int(max_miss):
                break
            continue
        x_new, e_new = pick
        prepend.append((float(t_prev), float(x_new)))
        inst_v = float((xb - x_new))
        vb = float(np.clip((0.60 * vb) + (0.40 * inst_v), -20.0, 20.0))
        xb = float(x_new)
        tb = int(t_prev)
        miss = 0
        steps += 1
        if e_new is not None:
            if eb is None:
                eb = e_new
            else:
                mix = (0.65 * eb) + (0.35 * e_new)
                n = float(np.linalg.norm(mix))
                if n > 1e-6:
                    eb = (mix / n).astype(np.float32, copy=False)

    if prepend:
        prepend.reverse()
    out = np.asarray(prepend + [(float(t), float(x)) for t, x in pb], dtype=np.float32)
    return _collapse_path_unique_t(out)


def _extend_tracks(
    tracks: list[dict],
    *,
    prob: np.ndarray,
    simplify_eps: float,
    max_anchors: int,
    emb_map: np.ndarray | None = None,
    max_extend_t: int = 18,
    search_dx: int = 6,
    min_prob: float = 0.05,
    max_miss: int = 2,
) -> list[dict]:
    if not tracks:
        return tracks
    out: list[dict] = []
    for tr in tracks:
        path = _track_path(tr)
        if path.shape[0] < 2:
            out.append(tr)
            continue
        ext = _extend_single_track_path(
            path,
            prob,
            emb_map=emb_map,
            max_extend_t=int(max_extend_t),
            search_dx=int(max(1, search_dx)),
            min_prob=float(max(0.0, min_prob)),
            max_miss=int(max(0, max_miss)),
        )
        rebuilt = _build_track_from_path(
            ext,
            [tr],
            simplify_eps=float(simplify_eps),
            max_anchors=int(max_anchors),
        )
        out.append(rebuilt if rebuilt is not None else tr)
    return out


def _track_overlap_stats(a: dict, b: dict) -> tuple[float, float, float, float, float, float, float]:
    pa = _collapse_path_unique_t(_track_path(a))
    pb = _collapse_path_unique_t(_track_path(b))
    if pa.shape[0] < 2 or pb.shape[0] < 2:
        return float("inf"), 0.0, float("inf"), float("inf"), float("inf"), 0.0, 0.0

    t0 = max(float(pa[0, 0]), float(pb[0, 0]))
    t1 = min(float(pa[-1, 0]), float(pb[-1, 0]))
    if t1 <= t0:
        return float("inf"), 0.0, float("inf"), float("inf"), float("inf"), 0.0, 0.0

    ts = np.arange(int(np.ceil(t0)), int(np.floor(t1)) + 1, dtype=np.float32)
    if ts.size < 3:
        return float("inf"), 0.0, float("inf"), float("inf"), float("inf"), 0.0, 0.0

    xa = np.interp(ts, pa[:, 0], pa[:, 1]).astype(np.float32)
    xb = np.interp(ts, pb[:, 0], pb[:, 1]).astype(np.float32)
    abs_dx = np.abs(xa - xb)
    mean_dx = float(np.mean(abs_dx))
    p90_dx = float(np.percentile(abs_dx, 90.0))
    start_dx = float(abs(float(xa[0]) - float(xb[0])))
    end_dx = float(abs(float(xa[-1]) - float(xb[-1])))
    overlap_frac = float(ts.size / max(1, min(pa.shape[0], pb.shape[0])))
    overlap_union = float(ts.size / max(1, (pa.shape[0] + pb.shape[0] - ts.size)))
    span_ratio = float(min(pa.shape[0], pb.shape[0]) / max(1, max(pa.shape[0], pb.shape[0])))
    return mean_dx, overlap_frac, p90_dx, start_dx, end_dx, overlap_union, span_ratio


def _dedupe_tracks(tracks: list[dict], max_mean_dx: float = 2.4, min_overlap_frac: float = 0.35) -> list[dict]:
    if len(tracks) < 2:
        return tracks

    ordered = sorted(
        tracks,
        key=lambda tr: (
            float(tr.get("score", 0.0)),
            float(tr.get("t_span", 0.0)),
            float(tr.get("length_px", 0.0)),
        ),
        reverse=True,
    )

    kept: list[dict] = []
    for tr in ordered:
        duplicate = False
        for ktr in kept:
            mean_dx, overlap, p90_dx, start_dx, end_dx, overlap_union, span_ratio = _track_overlap_stats(tr, ktr)
            if overlap < float(min_overlap_frac):
                continue

            pa = _collapse_path_unique_t(_track_path(tr))
            pb = _collapse_path_unique_t(_track_path(ktr))
            if pa.shape[0] < 2 or pb.shape[0] < 2:
                continue
            t0 = max(float(pa[0, 0]), float(pb[0, 0]))
            t1 = min(float(pa[-1, 0]), float(pb[-1, 0]))
            if t1 <= t0:
                continue
            ts = np.arange(int(np.ceil(t0)), int(np.floor(t1)) + 1, dtype=np.float32)
            if ts.size < 3:
                continue
            xa = np.interp(ts, pa[:, 0], pa[:, 1]).astype(np.float32)
            xb = np.interp(ts, pb[:, 0], pb[:, 1]).astype(np.float32)
            abs_dx = np.abs(xa - xb)
            close1 = float(np.mean(abs_dx <= float(max_mean_dx)))
            close2 = float(np.mean(abs_dx <= float(max(1.25, 1.8 * float(max_mean_dx)))))

            if mean_dx <= float(max_mean_dx) and close1 >= 0.62:
                duplicate = True
                break

            # Near-duplicates with minor local divergence still have strong
            # overlap + endpoint agreement + high close-frame fraction.
            strong_overlap = overlap >= float(max(0.78, min_overlap_frac))
            endpoint_tol = float(max(1.8, 2.6 * float(max_mean_dx)))
            body_tol = float(max(2.6, 3.4 * float(max_mean_dx)))
            if strong_overlap and close2 >= 0.74 and start_dx <= endpoint_tol and end_dx <= endpoint_tol and p90_dx <= body_tol:
                duplicate = True
                break

            # Containment case: short fragment riding on a longer track.
            # Use IoU-like overlap + high close fraction; this removes many
            # "same track but split into pieces" duplicates.
            contained = overlap_union >= 0.42 and span_ratio <= 0.72
            if contained and close2 >= 0.86 and mean_dx <= float(max(1.6, 2.4 * float(max_mean_dx))):
                duplicate = True
                break
        if not duplicate:
            kept.append(tr)
    return kept


def _extract_tracks_from_labels(
    labels: np.ndarray,
    n_labels: int,
    prob_work: np.ndarray,
    *,
    min_t_span: float,
    min_len: float,
    vmin: float,
    min_mean_prob: float,
    simplify_eps: float,
    max_anchors: int,
    min_directionality: float,
    min_net_velocity: float,
    emb_map: np.ndarray | None = None,
    max_row_span_median: float | None = None,
    max_row_span_p90: float | None = None,
    prefer_skeleton_dp: bool = False,
    debug_counts: dict | None = None,
    debug_reason_map: np.ndarray | None = None,
) -> list[dict]:
    out: list[dict] = []
    for comp_id in range(1, n_labels + 1):
        _bump_debug_count(debug_counts, "components_seen")
        comp = labels == comp_id
        if max_row_span_median is not None or max_row_span_p90 is not None:
            med_span, p90_span = _component_row_span_stats(comp)
            if max_row_span_median is not None and med_span > float(max_row_span_median):
                _bump_debug_count(debug_counts, "reject_row_span_median")
                _mark_reason_map(debug_reason_map, comp, "reject_row_span_median")
                continue
            if max_row_span_p90 is not None and p90_span > float(max_row_span_p90):
                _bump_debug_count(debug_counts, "reject_row_span_p90")
                _mark_reason_map(debug_reason_map, comp, "reject_row_span_p90")
                continue
        candidate_masks = [comp]
        if prefer_skeleton_dp:
            candidate_masks.extend(_skeleton_branch_candidate_masks(comp))

        def _build_track(candidate_mask: np.ndarray, *, count_rejects: bool) -> tuple[dict | None, str | None]:
            cand_rows = np.nonzero(np.any(candidate_mask, axis=1))[0]
            cand_row_count = int(cand_rows.size)
            used_path_rescue = False
            if prefer_skeleton_dp:
                path = _skeleton_label_dp_path(candidate_mask, prob=prob_work)
                if len(path) < 2:
                    path = _component_to_path(candidate_mask, prob=prob_work, emb_map=emb_map)
            else:
                path = _component_to_path(candidate_mask, prob=prob_work, emb_map=emb_map)
            if len(path) < 2:
                rescue_path = _skeleton_component_to_path(candidate_mask, prob=prob_work)
                if len(rescue_path) >= 2:
                    path = rescue_path
                    used_path_rescue = True
                    if count_rejects:
                        _bump_debug_count(debug_counts, "rescued_path_short")
                else:
                    if count_rejects:
                        _bump_debug_count(debug_counts, "reject_path_short")
                        if cand_row_count <= 1:
                            _bump_debug_count(debug_counts, "reject_path_short_single_row")
                        else:
                            _bump_debug_count(debug_counts, "reject_path_short_multi_row")
                    return None, "reject_path_short"

            t_span = float(path[-1][0] - path[0][0])
            if t_span < float(min_t_span):
                if count_rejects:
                    _bump_debug_count(debug_counts, "reject_t_span")
                return None, "reject_t_span"

            length = _polyline_len(path)
            if length < float(min_len):
                if count_rejects:
                    _bump_debug_count(debug_counts, "reject_length")
                return None, "reject_length"

            mean_abs_dxdt = _mean_abs_velocity(path)
            motion = _motion_strength(path)
            if motion < float(vmin):
                if count_rejects:
                    _bump_debug_count(debug_counts, "reject_vmin")
                return None, "reject_vmin"
            net_vel = _net_velocity(path)
            dir_score = _directionality(path)
            if float(min_directionality) > 0.0 and dir_score < float(min_directionality):
                if count_rejects:
                    _bump_debug_count(debug_counts, "reject_directionality")
                return None, "reject_directionality"
            if float(min_net_velocity) > 0.0 and net_vel < float(min_net_velocity):
                if count_rejects:
                    _bump_debug_count(debug_counts, "reject_net_velocity")
                return None, "reject_net_velocity"

            comp_mean_prob = float(np.mean(prob_work[candidate_mask])) if np.any(candidate_mask) else 0.0
            path_prob_samples = np.asarray(
                [_sample_prob_local_max(prob_work, int(round(t)), x, x_radius=1) for t, x in path],
                dtype=np.float32,
            )
            path_mean_prob = float(np.mean(path_prob_samples)) if path_prob_samples.size > 0 else comp_mean_prob
            path_p85_prob = (
                float(np.percentile(path_prob_samples, 85.0)) if path_prob_samples.size > 0 else path_mean_prob
            )
            track_conf = float(max(comp_mean_prob, 0.70 * path_mean_prob + 0.30 * path_p85_prob))
            if track_conf < float(min_mean_prob):
                if count_rejects:
                    _bump_debug_count(debug_counts, "reject_mean_prob")
                return None, "reject_mean_prob"

            simplified = _rdp(path, epsilon=float(simplify_eps))
            simplified = _cap_anchor_count(simplified, int(max_anchors))
            if len(simplified) < 2:
                if count_rejects:
                    _bump_debug_count(debug_counts, "reject_simplified_short")
                return None, "reject_simplified_short"

            anchors = [{"t": int(round(t)), "x": float(x)} for t, x in simplified]
            anchors.sort(key=lambda a: a["t"])
            score = float(track_conf * (length**0.72) * (1.0 + 0.18 * motion) * (1.0 + 0.008 * t_span))
            return {
                "anchors": anchors,
                "path": path,
                "length_px": float(length),
                "mean_abs_dxdt": float(mean_abs_dxdt),
                "motion_strength": float(motion),
                "net_velocity": float(net_vel),
                "directionality": float(dir_score),
                "mean_prob": float(track_conf),
                "component_mean_prob": float(comp_mean_prob),
                "path_mean_prob": float(path_mean_prob),
                "path_p85_prob": float(path_p85_prob),
                "t_span": float(t_span),
                "score": score,
                "path_source": "row_center_rescue" if used_path_rescue else ("skeleton_dp" if prefer_skeleton_dp else "ridge_trace"),
            }, None

        built_primary, primary_reject_reason = _build_track(comp, count_rejects=True)
        if built_primary is None:
            if primary_reject_reason is not None:
                _mark_reason_map(debug_reason_map, comp, primary_reject_reason)
            continue

        accepted_tracks = [built_primary]
        for extra_mask in candidate_masks[1:]:
            extra_track, _ = _build_track(extra_mask, count_rejects=False)
            if extra_track is not None:
                accepted_tracks.append(extra_track)

        for tr in accepted_tracks:
            out.append(tr)
            _bump_debug_count(debug_counts, "accepted")
        _mark_reason_map(debug_reason_map, comp, "accepted")
    return out


def postprocess_prob_to_tracks(
    prob_move: np.ndarray,
    prob_thresh: float = 0.35,
    min_len: float = 25.0,
    vmin: float = 0.15,
    simplify_eps: float = 2.0,
    max_anchors: int = 12,
    min_component_size: int = 18,
    min_t_span: int = 16,
    min_mean_prob: float = 0.48,
    max_tracks: int = 8,
    smooth_sigma_t: float = 0.6,
    smooth_sigma_x: float = 0.6,
    merge_gap_t: int = 24,
    merge_max_dx: float = 8.0,
    min_directionality: float = 0.10,
    min_net_velocity: float = 0.03,
    merge_min_bridge_prob: float | None = None,
    merge_min_bridge_coverage: float = 0.30,
    embedding_map: np.ndarray | None = None,
    merge_min_embed_sim: float | None = None,
    dedupe_max_mean_dx: float = 2.4,
    dedupe_min_overlap_frac: float = 0.35,
    keep_path: bool = False,
    debug: dict | None = None,
) -> list[dict]:
    prob = np.asarray(prob_move, dtype=np.float32)
    if prob.ndim != 2:
        raise ValueError(f"Expected (T, X) prob map, got shape {prob.shape!r}")

    emb_map: np.ndarray | None = None
    if embedding_map is not None:
        emb = np.asarray(embedding_map, dtype=np.float32)
        if emb.ndim != 3:
            raise ValueError(f"Expected embedding map with shape [D,T,X], got {emb.shape!r}")
        if emb.shape[1:] == prob.shape:
            emb_map = emb
        elif emb.shape[:2] == prob.shape:
            # Accept [T, X, D] and transpose.
            emb_map = np.transpose(emb, (2, 0, 1)).astype(np.float32, copy=False)
        else:
            raise ValueError(
                f"Embedding map spatial shape mismatch: prob={prob.shape!r}, embedding={emb.shape!r}"
            )

    prob_work = prob
    if smooth_sigma_t > 0 or smooth_sigma_x > 0:
        prob_work = ndi.gaussian_filter(prob_work, sigma=(max(0.0, float(smooth_sigma_t)), max(0.0, float(smooth_sigma_x))))
    if debug is not None:
        debug["prob_work"] = np.asarray(prob_work, dtype=np.float32)

    mask = _build_mask(
        prob_work,
        prob_thresh=float(prob_thresh),
        min_component_size=int(min_component_size),
        debug=debug,
    )
    labels, n = ndi.label(mask, structure=_STRUCTURE_8)
    if debug is not None:
        debug["mask_array"] = np.asarray(mask, dtype=np.uint8)
        debug["component_labels"] = np.asarray(labels, dtype=np.int32)
    component_rejects: dict = {}
    component_reason_map = np.zeros(labels.shape, dtype=np.uint8)
    scored_tracks = _extract_tracks_from_labels(
        labels,
        n,
        prob_work,
        min_t_span=float(min_t_span),
        min_len=float(min_len),
        vmin=float(vmin),
        min_mean_prob=float(min_mean_prob),
        simplify_eps=float(simplify_eps),
        max_anchors=int(max_anchors),
        min_directionality=float(min_directionality),
        min_net_velocity=float(min_net_velocity),
        emb_map=emb_map,
        # Filter broad merged blobs (common around overlap/split regions).
        max_row_span_median=4.0,
        max_row_span_p90=7.0,
        debug_counts=component_rejects,
        debug_reason_map=component_reason_map,
    )
    if debug is not None:
        debug["component_tracks"] = [_clone_track_debug(tr) for tr in scored_tracks]
        debug["component_rejects"] = dict(component_rejects)
        debug["component_reason_map"] = component_reason_map
        debug["reason_legend"] = {int(v): k for k, v in _REJECT_REASON_CODES.items()}

    if np.any(mask):
        skel = _skeletonize(mask)
        skel_labels, skel_n = ndi.label(skel, structure=_STRUCTURE_8)
        if debug is not None:
            debug["skeleton_array"] = np.asarray(skel, dtype=np.uint8)
            debug["skeleton_labels"] = np.asarray(skel_labels, dtype=np.int32)
        skeleton_rejects: dict = {}
        skeleton_reason_map = np.zeros(skel_labels.shape, dtype=np.uint8)
        skel_tracks = _extract_tracks_from_labels(
            skel_labels,
            skel_n,
            prob_work,
            min_t_span=float(min_t_span),
            min_len=float(min_len),
            vmin=float(vmin),
            min_mean_prob=float(min_mean_prob),
            simplify_eps=float(simplify_eps),
            max_anchors=int(max_anchors),
            min_directionality=float(min_directionality),
            min_net_velocity=float(min_net_velocity),
            emb_map=emb_map,
            max_row_span_median=None,
            max_row_span_p90=None,
            prefer_skeleton_dp=True,
            debug_counts=skeleton_rejects,
            debug_reason_map=skeleton_reason_map,
        )
        if skel_tracks:
            scored_tracks = _dedupe_tracks(
                scored_tracks + skel_tracks,
                max_mean_dx=float(max(0.1, dedupe_max_mean_dx)),
                min_overlap_frac=float(max(0.0, min(1.0, dedupe_min_overlap_frac))),
            )
        if debug is not None:
            debug["skeleton_tracks"] = [_clone_track_debug(tr) for tr in skel_tracks]
            debug["skeleton_rejects"] = dict(skeleton_rejects)
            debug["skeleton_reason_map"] = skeleton_reason_map
    if debug is not None:
        debug["premerge_tracks"] = [_clone_track_debug(tr) for tr in scored_tracks]

    if scored_tracks:
        bridge_prob_floor = (
            float(max(0.0, merge_min_bridge_prob))
            if merge_min_bridge_prob is not None
            else float(max(0.10, 0.45 * float(min_mean_prob)))
        )
        scored_tracks = _merge_split_tracks(
            scored_tracks,
            simplify_eps=float(simplify_eps),
            max_anchors=int(max_anchors),
            max_gap_t=int(max(1, merge_gap_t)),
            max_link_dx=float(max(1.0, merge_max_dx)),
            prob=prob_work,
            min_bridge_prob=bridge_prob_floor,
            min_bridge_coverage=float(max(0.0, min(1.0, merge_min_bridge_coverage))),
            emb_map=emb_map,
            min_embed_sim=merge_min_embed_sim,
        )
        scored_tracks = _dedupe_tracks(
            scored_tracks,
            max_mean_dx=float(max(0.1, dedupe_max_mean_dx)),
            min_overlap_frac=float(max(0.0, min(1.0, dedupe_min_overlap_frac))),
        )
        if debug is not None:
            debug["post_merge_tracks"] = [_clone_track_debug(tr) for tr in scored_tracks]

        # Follow likely ridge continuation beyond fragmented component boundaries.
        # Extension should only follow a clearly supported local continuation.
        # After the 8-connected labeling fix, most real tracks are already
        # extracted earlier; the extension pass should bridge short strong gaps,
        # not roam through weak stray pixels.
        ext_min_prob = float(
            np.clip(
                max(
                    0.18,
                    0.30 * float(prob_thresh),
                    0.50 * float(bridge_prob_floor),
                    0.35 * float(min_mean_prob),
                ),
                0.16,
                0.60,
            )
        )
        scored_tracks = _extend_tracks(
            scored_tracks,
            prob=prob_work,
            simplify_eps=float(simplify_eps),
            max_anchors=int(max_anchors),
            emb_map=emb_map,
            max_extend_t=int(max(4, min(14, int(max(1, merge_gap_t)) + 3))),
            search_dx=int(max(2, min(8, round(1.0 * float(max(1.0, merge_max_dx)))))),
            min_prob=ext_min_prob,
            max_miss=int(max(1, min(2, 1 + int(max(1, merge_gap_t)) // 6))),
        )
        scored_tracks = _merge_split_tracks(
            scored_tracks,
            simplify_eps=float(simplify_eps),
            max_anchors=int(max_anchors),
            max_gap_t=int(max(1, merge_gap_t)),
            max_link_dx=float(max(1.0, merge_max_dx)),
            prob=prob_work,
            min_bridge_prob=bridge_prob_floor,
            min_bridge_coverage=float(max(0.0, min(1.0, merge_min_bridge_coverage))),
            emb_map=emb_map,
            min_embed_sim=merge_min_embed_sim,
        )
        scored_tracks = _dedupe_tracks(
            scored_tracks,
            max_mean_dx=float(max(0.1, dedupe_max_mean_dx)),
            min_overlap_frac=float(max(0.0, min(1.0, dedupe_min_overlap_frac))),
        )
        if debug is not None:
            debug["post_extend_tracks"] = [_clone_track_debug(tr) for tr in scored_tracks]

    if max_tracks and max_tracks > 0 and len(scored_tracks) > max_tracks:
        scored_tracks.sort(key=lambda tr: tr.get("score", 0.0), reverse=True)
        scored_tracks = scored_tracks[: int(max_tracks)]

    tracks = []
    for tr in scored_tracks:
        clean = dict(tr)
        if not keep_path:
            clean.pop("path", None)
        tracks.append(clean)
    tracks.sort(key=lambda tr: tr["anchors"][0]["t"] if tr["anchors"] else 10**9)
    if debug is not None:
        debug["final_tracks"] = [_clone_track_debug(tr) for tr in tracks]
        debug["counts"] = {
            "component": len(debug.get("component_tracks", [])),
            "skeleton": len(debug.get("skeleton_tracks", [])),
            "premerge": len(debug.get("premerge_tracks", [])),
            "post_merge": len(debug.get("post_merge_tracks", [])),
            "post_extend": len(debug.get("post_extend_tracks", [])),
            "final": len(debug.get("final_tracks", [])),
        }
    return tracks


def tracks_to_json(tracks: list[dict]) -> str:
    return json.dumps({"tracks": tracks}, indent=2)
