"""Person tracking and bad-detection removal for RTMPose multi-person output.

Ported from ``rtm_keypoints_detector/filtering.py``.  Pure numpy — no GPU
dependencies.
"""

from __future__ import annotations

import numpy as np


def compute_person_correspondences(
    kp_list: list[np.ndarray],
    bboxes_list: list[np.ndarray],
    max_distance: float = 150,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Align person detections across frames using distance-based matching.

    Uses multiple keypoints and bounding box centres for robust matching.

    Args:
        kp_list: Per-frame keypoint arrays ``[(n_people_i, 133, 3), ...]``.
        bboxes_list: Per-frame bbox arrays ``[(n_people_i, 2, 3), ...]``.
        max_distance: Maximum distance threshold for matching.

    Returns:
        Tuple of ``(aligned_keypoints, aligned_bboxes, person_masks)``
        where shapes are ``(T, N, 133, 3)``, ``(T, N, 2, 3)``, ``(T, N)``.
        NaN fills missing detections.
    """
    if not kp_list or all(len(kp) == 0 for kp in kp_list):
        return np.array([]), np.array([]), np.array([])

    max_people = max(len(kp) for kp in kp_list if len(kp) > 0)
    n_frames = len(kp_list)

    # Find first non-empty frame
    first_frame_idx = 0
    while first_frame_idx < n_frames and len(kp_list[first_frame_idx]) == 0:
        first_frame_idx += 1

    if first_frame_idx >= n_frames:
        return np.array([]), np.array([]), np.array([])

    dtype = kp_list[first_frame_idx].dtype
    aligned_keypoints = np.full(
        (n_frames, max_people, 133, 3), np.nan, dtype=dtype,
    )
    aligned_bboxes = np.full(
        (n_frames, max_people, 2, 3), np.nan, dtype=dtype,
    )
    person_masks = np.zeros((n_frames, max_people), dtype=bool)

    if max_people == 0:
        return np.array([]), np.array([]), np.array([])

    # Seed with first frame
    first_kp = kp_list[first_frame_idx]
    first_bbox = bboxes_list[first_frame_idx]
    n_first = len(first_kp)
    aligned_keypoints[first_frame_idx, :n_first] = first_kp
    aligned_bboxes[first_frame_idx, :n_first] = first_bbox
    person_masks[first_frame_idx, :n_first] = True

    def _compute_person_features(
        keypoints: np.ndarray,
        bboxes: np.ndarray,
    ) -> np.ndarray:
        features = []
        for i in range(len(keypoints)):
            kp = keypoints[i]
            bbox = bboxes[i]
            bbox_center = (bbox[0] + bbox[1]) / 2

            valid_points = []
            for pt in [kp[5], kp[6], kp[11], kp[12]]:  # shoulders + hips
                if not np.any(np.isnan(pt)):
                    valid_points.append(pt)

            torso_center = (
                np.mean(valid_points, axis=0) if valid_points else bbox_center
            )

            if not np.any(np.isnan(torso_center)):
                feature = 0.3 * bbox_center + 0.7 * torso_center
            else:
                feature = bbox_center

            features.append(feature)
        return np.array(features)

    def _solve_assignment(
        cost_matrix: np.ndarray,
        max_cost: float,
    ) -> dict[int, int]:
        n_current, n_prev = cost_matrix.shape
        assignments: dict[int, int] = {}
        used_prev: set[int] = set()

        candidates = []
        for i in range(n_current):
            for j in range(n_prev):
                if cost_matrix[i, j] <= max_cost:
                    candidates.append((cost_matrix[i, j], i, j))
        candidates.sort()

        for _cost, curr_idx, prev_idx in candidates:
            if curr_idx not in assignments and prev_idx not in used_prev:
                assignments[curr_idx] = prev_idx
                used_prev.add(prev_idx)
        return assignments

    for frame_idx in range(n_frames):
        if frame_idx == first_frame_idx:
            continue

        current_kp = kp_list[frame_idx]
        current_bbox = bboxes_list[frame_idx]
        if len(current_kp) == 0:
            continue

        # Most recent frame with detections
        ref_frame_idx = frame_idx - 1
        while ref_frame_idx >= 0 and not np.any(person_masks[ref_frame_idx]):
            ref_frame_idx -= 1

        if ref_frame_idx < 0:
            n_current = len(current_kp)
            aligned_keypoints[frame_idx, :n_current] = current_kp
            aligned_bboxes[frame_idx, :n_current] = current_bbox
            person_masks[frame_idx, :n_current] = True
            continue

        prev_valid_mask = person_masks[ref_frame_idx]
        prev_indices = np.where(prev_valid_mask)[0]

        if len(prev_indices) == 0:
            n_current = len(current_kp)
            aligned_keypoints[frame_idx, :n_current] = current_kp
            aligned_bboxes[frame_idx, :n_current] = current_bbox
            person_masks[frame_idx, :n_current] = True
            continue

        prev_kp = aligned_keypoints[ref_frame_idx, prev_indices]
        prev_bbox = aligned_bboxes[ref_frame_idx, prev_indices]

        current_features = _compute_person_features(current_kp, current_bbox)
        prev_features = _compute_person_features(prev_kp, prev_bbox)

        cost_matrix = np.full(
            (len(current_kp), len(prev_indices)), np.inf,
        )
        for i, curr_feat in enumerate(current_features):
            for j, prev_feat in enumerate(prev_features):
                if not (np.any(np.isnan(curr_feat)) or np.any(np.isnan(prev_feat))):
                    cost_matrix[i, j] = np.linalg.norm(curr_feat - prev_feat)

        assignments = _solve_assignment(cost_matrix, max_distance)

        assigned_current: set[int] = set()
        for curr_idx, prev_local_idx in assignments.items():
            prev_global_idx = prev_indices[prev_local_idx]
            aligned_keypoints[frame_idx, prev_global_idx] = current_kp[curr_idx]
            aligned_bboxes[frame_idx, prev_global_idx] = current_bbox[curr_idx]
            person_masks[frame_idx, prev_global_idx] = True
            assigned_current.add(curr_idx)

        unassigned_current = set(range(len(current_kp))) - assigned_current
        available_slots = [
            s for s in range(max_people) if not person_masks[frame_idx, s]
        ]

        for i, curr_idx in enumerate(
            list(unassigned_current)[: len(available_slots)]
        ):
            slot_idx = available_slots[i]
            aligned_keypoints[frame_idx, slot_idx] = current_kp[curr_idx]
            aligned_bboxes[frame_idx, slot_idx] = current_bbox[curr_idx]
            person_masks[frame_idx, slot_idx] = True

    return aligned_keypoints, aligned_bboxes, person_masks


def remove_bad_detections(
    aligned_kp: np.ndarray,
    aligned_bboxes: np.ndarray,
    person_masks: np.ndarray,
    frame_shape: tuple[int, int],
    bbox_area_thresh: float = 0.1,
    nan_detections_thresh: float = 0.9,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Remove noisy / spurious person detections.

    Removes persons whose detections are mostly NaN or whose bounding boxes
    are consistently too small or span the full frame.

    Args:
        aligned_kp: ``(T, N, 133, 3)`` aligned keypoints.
        aligned_bboxes: ``(T, N, 2, 3)`` aligned bounding boxes.
        person_masks: ``(T, N)`` boolean masks.
        frame_shape: ``(height, width)`` of the video frames.
        bbox_area_thresh: Minimum bbox-area / frame-area ratio.
        nan_detections_thresh: Maximum fraction of NaN frames per person.

    Returns:
        Cleaned ``(aligned_kp, aligned_bboxes, person_masks)`` with bad
        persons removed along axis 1.
    """
    n_frames, n_people, _, _ = aligned_kp.shape

    full_frame_bbox = np.array(
        [[0.0, 0.0], [frame_shape[1], frame_shape[0]]],
    )

    # Full-frame bboxes with confidence 1.0 are usually bad detections
    bad_full_frame_bbox_mask = np.logical_and(
        (aligned_bboxes[:, :, :, :2] == full_frame_bbox).all(axis=3).all(axis=2),
        (aligned_bboxes[:, :, :, 2] == 1.0).all(axis=2),
    )
    aligned_kp[bad_full_frame_bbox_mask] = np.nan
    aligned_bboxes[bad_full_frame_bbox_mask] = np.nan
    person_masks[bad_full_frame_bbox_mask] = False

    # Person is NaN for >90% of frames → remove
    n_nan = (
        np.isnan(aligned_kp[:, :, :, :2]).all(axis=3).all(axis=2).sum(axis=0)
    )
    all_nan_detections = (n_nan / n_frames) > nan_detections_thresh

    # Bbox area < threshold for all valid frames → remove
    frame_area = frame_shape[0] * frame_shape[1]
    x1, y1 = aligned_bboxes[:, :, 0, :2].transpose(2, 0, 1).astype(np.float32)
    x2, y2 = aligned_bboxes[:, :, 1, :2].transpose(2, 0, 1).astype(np.float32)
    bboxes_area = abs(x2 - x1) * abs(y2 - y1)

    all_small_bboxes = np.array([False for _ in range(n_people)])
    for pi in range(n_people):
        p_area = bboxes_area[:, pi]
        valid_area = p_area[~np.isnan(p_area)]
        if len(valid_area) > 0 and ((valid_area / frame_area) < bbox_area_thresh).all():
            all_small_bboxes[pi] = True

    to_remove = np.logical_or(all_nan_detections, all_small_bboxes)
    if to_remove.any():
        return (
            aligned_kp[:, ~to_remove],
            aligned_bboxes[:, ~to_remove],
            person_masks[:, ~to_remove],
        )

    return aligned_kp, aligned_bboxes, person_masks
