"""Resident native TensorRT inference with causal pose fusion and tracking."""
from __future__ import annotations
import time
from collections import defaultdict, deque
import numpy as np
from pose_tracking import BoxTracker

SIDES = ("left", "right")

class NativePoseEngine:
    def __init__(self, args, native):
        self.args = args
        self.native = native
        self.pose = native.pose
        self.box_tracker = BoxTracker(max_age=30)
        self.people = {}
        self.timings = defaultdict(lambda: deque(maxlen=3000))
        self.isolated_inferences = 0

    def infer(self, frame: np.ndarray, frame_idx: int) -> dict:
        from pose_geometry import HAND_SLICES
        from pose_fusion import adaptive_smooth, build_global_hands, fuse_person
        from pose_occlusion import Measurement, resolve_global_hand_assignments
        from pose_tracking import (
            LivePersonState,
            mask_other_people,
            update_live_v2,
        )

        infer_started = time.perf_counter()
        height, width = frame.shape[:2]
        started = time.perf_counter()
        detection_height = max(
            1, min(height, round(height * self.args.det_roi_y_max_ratio))
        )
        boxes, confidences = self.native.detect(frame[:detection_height])
        self.timings["detector_ms"].append((time.perf_counter() - started) * 1000)
        track_ids = self.box_tracker.update(boxes, frame_idx, (height, width))

        # Retire disappeared people in long-running camera sessions.
        for stale_id in list(self.people):
            if stale_id not in self.box_tracker.tracks:
                del self.people[stale_id]
        if boxes:
            started = time.perf_counter()
            raw_keypoints, raw_scores = self.pose(frame, bboxes=boxes)
            self.timings["raw_pose_ms"].append((time.perf_counter() - started) * 1000)
            isolated_keypoints = raw_keypoints.copy()
            isolated_scores = raw_scores.copy()
            started = time.perf_counter()
            for person_idx, bbox in enumerate(boxes):
                target_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0)
                needs_isolation = any(
                    idx != person_idx
                    and max((other[2] - other[0]) * (other[3] - other[1]), 1.0)
                    < 0.5 * target_area
                    for idx, other in enumerate(boxes)
                )
                if not needs_isolation:
                    continue
                masked = mask_other_people(frame, boxes, person_idx)
                keypoints, scores = self.pose(masked, bboxes=[bbox])
                isolated_keypoints[person_idx] = keypoints[0]
                isolated_scores[person_idx] = scores[0]
                self.isolated_inferences += 1
            self.timings["isolated_extra_ms"].append(
                (time.perf_counter() - started) * 1000
            )
        else:
            raw_keypoints = np.empty((0, 133, 2), dtype=np.float64)
            raw_scores = np.empty((0, 133), dtype=np.float64)
            isolated_keypoints = raw_keypoints.copy()
            isolated_scores = raw_scores.copy()

        selected_list, smoothed_list, metadata_list, states = [], [], [], []
        for person_idx, (bbox, track_id) in enumerate(zip(boxes, track_ids)):
            state = self.people.setdefault(track_id, LivePersonState())
            selected, metadata = fuse_person(
                raw_keypoints[person_idx],
                raw_scores[person_idx],
                isolated_keypoints[person_idx],
                isolated_scores[person_idx],
                bbox,
                state.fusion,
            )
            selected_list.append(selected)
            smoothed_list.append(adaptive_smooth(selected, bbox, state.fusion))
            metadata_list.append(metadata)
            states.append(state)

        measurements, endpoint_tracks = {}, {}
        for person_idx, (selected, metadata, state) in enumerate(
            zip(selected_list, metadata_list, states)
        ):
            for side, (start, end) in HAND_SLICES.items():
                meta = metadata[side]
                measurements[(person_idx, side)] = (
                    Measurement(
                        keypoints=selected[start:end],
                        precise=bool(meta["distance_eligible"]),
                        provisional=not bool(meta["distance_eligible"]),
                        quality=float(meta.get("selected_hand_score", 0.0)),
                    )
                    if meta["alarm_eligible"]
                    else None
                )
                endpoint_tracks[(person_idx, side)] = state.endpoints[side]
        measurements, assignments = resolve_global_hand_assignments(
            measurements, endpoint_tracks, boxes, 30
        )

        v2_list, v2_meta_list = [], []
        for person_idx, (bbox, selected, smoothed, metadata, state) in enumerate(
            zip(boxes, selected_list, smoothed_list, metadata_list, states)
        ):
            v2, v2_meta = update_live_v2(
                frame_idx,
                bbox,
                selected,
                smoothed,
                metadata,
                state,
                {side: measurements[(person_idx, side)] for side in SIDES},
                {side: assignments.get((person_idx, side)) for side in SIDES},
                12,
                30,
            )
            v2_list.append(v2)
            v2_meta_list.append(v2_meta)

        self.timings["pose_pipeline_total_ms"].append(
            (time.perf_counter() - infer_started) * 1000
        )
        return {
            "frame_index": frame_idx,
            "track_ids": track_ids,
            "boxes": np.round(np.asarray(boxes), 3).tolist() if boxes else [],
            "det_confidences": [round(value, 4) for value in confidences],
            "raw_keypoints": np.round(raw_keypoints, 3).tolist(),
            "temporal_selected_keypoints": np.round(
                np.asarray(selected_list), 3
            ).tolist(),
            "temporal_keypoints": np.round(np.asarray(smoothed_list), 3).tolist(),
            "temporal_metadata": metadata_list,
            "v2_keypoints": np.round(np.asarray(v2_list), 3).tolist(),
            "v2_metadata": v2_meta_list,
            "global_hands": build_global_hands(selected_list, metadata_list),
        }