"""RTMPose whole-body keypoint processor.

Wraps MMPose's ``MMPoseInferencer`` with the ``rtmw-x`` model to extract
133 COCO-WholeBody keypoints per person per frame.

Ported from ``rtm_keypoints_detector/extract_keypoints_core.py``.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import Any

import numpy as np

logger = logging.getLogger(__name__)


@dataclass
class RTMPoseProcessorConfig:
    """Configuration for :class:`RTMPoseProcessor`."""

    model_name: str = "rtmw-x_8xb320-270e_cocktail14-384x288"
    kpt_thr: float = 0.5
    bbox_thr: float = 0.5
    kp_dtype: type = np.float16


class RTMPoseProcessor:
    """Bundles MMPose inference for whole-body keypoint extraction.

    Usage::

        cfg = RTMPoseProcessorConfig()
        proc = RTMPoseProcessor(cfg, device="cuda:0")
        result = proc.process_frames(frames, frame_indices)

    Args:
        config: Processor configuration.
        device: Torch device string (e.g. ``"cuda:0"``).
    """

    def __init__(self, config: RTMPoseProcessorConfig, device: str = "cuda:0") -> None:
        self.config = config
        self.device = device

        from mmpose.apis import MMPoseInferencer

        self._inferencer = MMPoseInferencer(config.model_name, device=device)
        logger.info(
            "RTMPose inferencer loaded: model=%s, device=%s",
            config.model_name,
            device,
        )

    def process_frames(
        self,
        frames: list[np.ndarray],
        frame_indices: list[int],
    ) -> dict[str, Any]:
        """Run MMPose inference on a batch of frames.

        Args:
            frames: List of ``(H, W, 3)`` uint8 numpy arrays (RGB or BGR —
                MMPose handles conversion internally).
            frame_indices: Corresponding absolute frame indices.

        Returns:
            Dict with per-frame detection lists:
              - ``all_kp``: list of ``(n_people, 133, 3)`` arrays per frame
              - ``all_bboxes``: list of ``(n_people, 2, 3)`` arrays per frame
              - ``all_bboxes_scores``: list of score lists per frame
              - ``successful_frames_per_person``: ``{person_id: [frame_indices]}``
              - ``num_total_detections``: total count
        """
        cfg = self.config
        kp_dtype = cfg.kp_dtype

        all_kp: list[np.ndarray | list] = []
        all_bboxes: list[np.ndarray | list] = []
        all_bboxes_scores: list[list[float] | list] = []
        successful_frames_per_person: dict[int, list[int]] = {}
        num_total_detections = 0

        result_generator = self._inferencer(
            frames,
            show=False,
            kpt_thr=cfg.kpt_thr,
            bbox_thr=cfg.bbox_thr,
            return_vis=False,
        )

        for sample_i in range(len(frames)):
            try:
                sample_result = next(result_generator)
            except StopIteration:
                all_kp.append([])
                all_bboxes.append([])
                all_bboxes_scores.append([])
                continue

            people_preds = sample_result["predictions"][0]
            n_det = len(people_preds)
            num_total_detections += n_det

            if n_det > 0:
                frame_kp = np.zeros((n_det, 133, 3), dtype=kp_dtype)
                frame_bbox = np.zeros((n_det, 2, 3), dtype=kp_dtype)
                scores: list[float] = []

                for person_i, pred in enumerate(people_preds):
                    frame_kp[person_i, :, :2] = np.array(
                        pred["keypoints"], dtype=kp_dtype,
                    )
                    frame_kp[person_i, :, 2] = np.array(
                        pred["keypoint_scores"], dtype=kp_dtype,
                    )

                    frame_bbox[person_i, :, :2] = np.array(
                        pred["bbox"][0], dtype=kp_dtype,
                    ).reshape(2, 2)
                    frame_bbox[person_i, :, 2] = np.expand_dims(
                        np.array(pred["bbox_score"], dtype=kp_dtype), axis=0,
                    ).repeat(2, axis=0)

                    if person_i not in successful_frames_per_person:
                        successful_frames_per_person[person_i] = []
                    successful_frames_per_person[person_i].append(
                        frame_indices[sample_i],
                    )
                    scores.append(pred["bbox_score"])

                all_kp.append(frame_kp)
                all_bboxes.append(frame_bbox)
                all_bboxes_scores.append(scores)
            else:
                all_kp.append([])
                all_bboxes.append([])
                all_bboxes_scores.append([])

        return {
            "all_kp": all_kp,
            "all_bboxes": all_bboxes,
            "all_bboxes_scores": all_bboxes_scores,
            "successful_frames_per_person": successful_frames_per_person,
            "num_total_detections": num_total_detections,
        }
