"""Runtime pose tracking helpers."""
from __future__ import annotations
import math
from collections import deque
from dataclasses import dataclass, field
import numpy as np
from pose_geometry import HAND_SLICES
from pose_fusion import BODY_SIDE_INDICES, PersonTrack
from pose_occlusion import EndpointTrack, Measurement, calibrated_bone_lengths, reachable_region, solve_two_bone_elbow, temporal_recovery_measurement, update_endpoint_track


def bbox_iou(a: np.ndarray, b: np.ndarray) -> float:
    x1 = max(float(a[0]), float(b[0]))
    y1 = max(float(a[1]), float(b[1]))
    x2 = min(float(a[2]), float(b[2]))
    y2 = min(float(a[3]), float(b[3]))
    intersection = max(x2 - x1, 0.0) * max(y2 - y1, 0.0)
    union = (
        max(float(a[2] - a[0]), 0.0) * max(float(a[3] - a[1]), 0.0)
        + max(float(b[2] - b[0]), 0.0) * max(float(b[3] - b[1]), 0.0)
        - intersection
    )
    return intersection / max(union, 1.0)


class BoxTracker:
    def __init__(self, max_age: int = 30):
        self.max_age = max_age
        self.next_id = 1
        self.tracks: dict[int, dict] = {}

    def update(
        self,
        boxes: list[list[float]],
        frame_idx: int,
        image_shape: tuple[int, int],
    ) -> list[int]:
        height, width = image_shape
        image_diagonal = math.hypot(width, height)
        active_ids = [
            track_id
            for track_id, track in self.tracks.items()
            if frame_idx - track["last_frame"] <= self.max_age
        ]
        self.tracks = {track_id: self.tracks[track_id] for track_id in active_ids}
        pairs = []
        for detection_idx, box in enumerate(boxes):
            box_array = np.asarray(box, dtype=np.float64)
            center = (box_array[:2] + box_array[2:]) * 0.5
            for track_id in active_ids:
                prior = np.asarray(self.tracks[track_id]["box"], dtype=np.float64)
                prior_center = (prior[:2] + prior[2:]) * 0.5
                iou = bbox_iou(box_array, prior)
                center_distance = float(np.linalg.norm(center - prior_center)) / image_diagonal
                if iou >= 0.05 or center_distance <= 0.18:
                    score = 2.0 * iou + max(0.0, 1.0 - center_distance / 0.18)
                    pairs.append((score, detection_idx, track_id))

        assigned_detections = set()
        assigned_tracks = set()
        result: list[int | None] = [None] * len(boxes)
        for _, detection_idx, track_id in sorted(pairs, reverse=True):
            if detection_idx in assigned_detections or track_id in assigned_tracks:
                continue
            assigned_detections.add(detection_idx)
            assigned_tracks.add(track_id)
            result[detection_idx] = track_id

        for detection_idx, box in enumerate(boxes):
            track_id = result[detection_idx]
            if track_id is None:
                track_id = self.next_id
                self.next_id += 1
                result[detection_idx] = track_id
            self.tracks[track_id] = {"box": box, "last_frame": frame_idx}
        return [int(track_id) for track_id in result]


@dataclass
class LivePersonState:
    fusion: PersonTrack = field(default_factory=PersonTrack)
    endpoints: dict[str, EndpointTrack] = field(
        default_factory=lambda: {"left": EndpointTrack(), "right": EndpointTrack()}
    )
    bone_samples: dict[str, deque] = field(
        default_factory=lambda: {
            "left": deque(maxlen=120),
            "right": deque(maxlen=120),
        }
    )
    last_frame: int = -1


def bbox_diagonal(bbox: list[float]) -> float:
    return max(math.hypot(bbox[2] - bbox[0], bbox[3] - bbox[1]), 1.0)


def update_bone_samples(
    state: LivePersonState,
    keypoints: np.ndarray,
    metadata: dict,
    bbox: list[float],
) -> None:
    diagonal = bbox_diagonal(bbox)
    for side, (shoulder, elbow, wrist) in BODY_SIDE_INDICES.items():
        meta = metadata[side]
        if meta["limb_state"] != "consensus" or not meta["distance_eligible"]:
            continue
        upper = float(np.linalg.norm(keypoints[shoulder] - keypoints[elbow])) / diagonal
        forearm = float(np.linalg.norm(keypoints[elbow] - keypoints[wrist])) / diagonal
        if 0.04 <= upper <= 0.65 and 0.04 <= forearm <= 0.65:
            state.bone_samples[side].append((upper, forearm))


def current_bone_ratios(state: LivePersonState, side: str) -> tuple[float, float]:
    samples = state.bone_samples[side]
    if len(samples) < 10:
        opposite = "right" if side == "left" else "left"
        if len(state.bone_samples[opposite]) >= 10:
            samples = state.bone_samples[opposite]
    if not samples:
        return 0.24, 0.22
    values = np.asarray(samples, dtype=np.float64)
    return float(np.median(values[:, 0])), float(np.median(values[:, 1]))


def update_live_v2(
    frame_idx: int,
    bbox: list[float],
    selected_keypoints: np.ndarray,
    smoothed_keypoints: np.ndarray,
    fusion_metadata: dict,
    state: LivePersonState,
    assigned_measurements: dict[str, Measurement | None] | None,
    identity_assignments: dict[str, dict] | None,
    short_occlusion_frames: int,
    medium_occlusion_frames: int,
) -> tuple[np.ndarray, dict]:
    frame_gap = 0 if state.last_frame < 0 else max(frame_idx - state.last_frame - 1, 0)
    for endpoint in state.endpoints.values():
        endpoint.missed_frames += frame_gap
    state.last_frame = frame_idx
    update_bone_samples(state, selected_keypoints, fusion_metadata, bbox)

    output = smoothed_keypoints.copy()
    output_metadata = {}
    diagonal = bbox_diagonal(bbox)
    for side, (shoulder_idx, elbow_idx, wrist_idx) in BODY_SIDE_INDICES.items():
        endpoint = state.endpoints[side]
        meta = fusion_metadata[side]
        start, end = HAND_SLICES[side]
        measurement = (
            assigned_measurements.get(side)
            if assigned_measurements is not None
            else None
        )
        if assigned_measurements is None and meta["alarm_eligible"]:
            measurement = Measurement(
                keypoints=selected_keypoints[start:end],
                precise=bool(meta["distance_eligible"]),
                provisional=not bool(meta["distance_eligible"]),
                quality=float(meta.get("selected_hand_score", 0.0)),
            )
        if measurement is None:
            measurement = temporal_recovery_measurement(
                selected_keypoints,
                meta,
                side,
                bbox,
                endpoint,
                medium_occlusion_frames,
            )
        returning_after_long = (
            measurement is not None
            and endpoint.missed_frames > short_occlusion_frames
        )
        wrist, hand, endpoint_state = update_endpoint_track(
            endpoint,
            measurement,
            frame_idx,
            short_occlusion_frames,
            medium_occlusion_frames,
        )
        shoulder = smoothed_keypoints[shoulder_idx]
        upper_ratio, forearm_ratio = current_bone_ratios(state, side)
        upper_length, forearm_length = calibrated_bone_lengths(
            upper_ratio, forearm_ratio, bbox
        )
        region = None
        if wrist is not None and hand is not None:
            reference_elbow = (
                selected_keypoints[elbow_idx]
                if returning_after_long or endpoint.last_elbow is None
                else endpoint.last_elbow
            )
            use_ik = (
                meta["limb_state"] != "consensus"
                or endpoint_state
                in ("provisional", "recovered", "reassigned", "predicted_short")
            )
            elbow = (
                solve_two_bone_elbow(
                    shoulder,
                    wrist,
                    upper_length,
                    forearm_length,
                    reference_elbow,
                )
                if use_ik
                else smoothed_keypoints[elbow_idx]
            )
            endpoint.last_elbow = elbow.copy()
            direction = wrist - shoulder
            if float(np.linalg.norm(direction)) > 1e-6:
                endpoint.last_direction = direction / np.linalg.norm(direction)
            output[shoulder_idx] = shoulder
            output[elbow_idx] = elbow
            output[wrist_idx] = wrist
            output[start:end] = hand
            uncertainty = (
                0.025 * diagonal
                if endpoint_state == "observed"
                else (0.04 + 0.015 * endpoint.missed_frames) * diagonal
            )
            precise = endpoint_state == "observed"
        else:
            region = reachable_region(
                shoulder,
                upper_length,
                forearm_length,
                bbox,
                endpoint,
                endpoint_state,
                short_occlusion_frames,
                medium_occlusion_frames,
            )
            uncertainty = region["outer_radius_px"]
            precise = False
        output_metadata[side] = {
            "state": endpoint_state,
            "precise_distance": precise,
            "fast_alarm_candidate": True,
            "risk_mode": (
                "precise_hand"
                if endpoint_state == "observed"
                else (
                    "motion_corridor_intersection"
                    if endpoint_state == "reachable_medium"
                    else (
                        "door_intersection_only"
                        if endpoint_state == "reachable_long"
                        else (
                            "reassigned_hand"
                            if endpoint_state == "reassigned"
                            else "temporal_hand"
                        )
                    )
                )
            ),
            "requires_door_intersection": endpoint_state
            in ("reachable_medium", "reachable_long"),
            "identity_assignment": (
                identity_assignments.get(side)
                if identity_assignments is not None
                else None
            ),
            "missed_frames": endpoint.missed_frames,
            "uncertainty_radius_px": round(float(uncertainty), 3),
            "bone_lengths_px": {
                "upper": round(upper_length, 3),
                "forearm": round(forearm_length, 3),
            },
            "reachable_region": region,
        }
    return output, output_metadata


def mask_other_people(
    frame: np.ndarray,
    boxes: list[list[float]],
    target_idx: int,
) -> np.ndarray:
    masked = frame.copy()
    height, width = frame.shape[:2]
    mask = np.zeros((height, width), dtype=bool)
    target = boxes[target_idx]
    target_area = max((target[2] - target[0]) * (target[3] - target[1]), 1.0)
    for idx, bbox in enumerate(boxes):
        if idx == target_idx:
            continue
        other_area = max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0)
        # The failure in this video is the small background person being
        # absorbed by the much larger foreground person's 3:4 pose crop.
        # Mask only markedly smaller people.  Masking a larger occluder while
        # estimating a small person removes genuine target pixels at overlap.
        if other_area >= 0.5 * target_area:
            continue
        x1, y1, x2, y2 = (int(round(value)) for value in bbox)
        x1, x2 = max(0, x1), min(width - 1, x2)
        y1, y2 = max(0, y1), min(height - 1, y2)
        mask[y1:y2 + 1, x1:x2 + 1] = True

    masked[mask] = (124, 116, 104)
    return masked