"""Runtime pose occlusion helpers."""
from __future__ import annotations
import math
from dataclasses import dataclass
import numpy as np
from pose_geometry import HAND_SLICES


@dataclass
class Measurement:
    keypoints: np.ndarray
    precise: bool
    provisional: bool
    source: str = "model"
    quality: float = 0.0


@dataclass
class EndpointTrack:
    wrist: np.ndarray | None = None
    velocity: np.ndarray | None = None
    hand_shape: np.ndarray | None = None
    last_observed_frame: int | None = None
    missed_frames: int = 0
    last_elbow: np.ndarray | None = None
    last_direction: np.ndarray | None = None


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


def solve_two_bone_elbow(
    shoulder: np.ndarray,
    wrist: np.ndarray,
    upper_length: float,
    forearm_length: float,
    reference_elbow: np.ndarray,
) -> np.ndarray:
    direction = wrist - shoulder
    distance = float(np.linalg.norm(direction))
    if distance < 1e-6:
        return reference_elbow.copy()
    unit = direction / distance

    if distance >= upper_length + forearm_length:
        return shoulder + unit * upper_length
    if distance <= abs(upper_length - forearm_length):
        sign = 1.0 if upper_length >= forearm_length else -1.0
        return shoulder + unit * upper_length * sign

    along = (
        upper_length * upper_length
        - forearm_length * forearm_length
        + distance * distance
    ) / (2.0 * distance)
    height = math.sqrt(max(upper_length * upper_length - along * along, 0.0))
    base = shoulder + unit * along
    perpendicular = np.asarray([-unit[1], unit[0]])
    candidates = (base + perpendicular * height, base - perpendicular * height)
    return min(candidates, key=lambda candidate: float(np.linalg.norm(candidate - reference_elbow)))


def update_endpoint_track(
    track: EndpointTrack,
    measurement: Measurement | None,
    frame_idx: int,
    short_occlusion_frames: int,
    medium_occlusion_frames: int,
) -> tuple[np.ndarray | None, np.ndarray | None, str]:
    if measurement is not None:
        measured_wrist = measurement.keypoints[0]
        measured_shape = measurement.keypoints - measured_wrist
        if track.wrist is None or track.missed_frames > short_occlusion_frames:
            track.wrist = measured_wrist.copy()
            track.velocity = np.zeros(2, dtype=np.float64)
            track.hand_shape = measured_shape.copy()
        else:
            prediction = track.wrist + track.velocity
            residual = measured_wrist - prediction
            alpha = 0.78 if measurement.precise else 0.88
            beta = 0.18 if measurement.precise else 0.24
            track.wrist = prediction + alpha * residual
            track.velocity = 0.82 * track.velocity + beta * residual
            shape_alpha = 0.55 if measurement.precise else 0.7
            track.hand_shape = (
                shape_alpha * measured_shape + (1.0 - shape_alpha) * track.hand_shape
            )
        track.last_observed_frame = frame_idx
        track.missed_frames = 0
        if measurement.source == "global_reassignment":
            state = "reassigned"
        elif measurement.source == "temporal_recovery":
            state = "recovered"
        elif measurement.precise:
            state = "observed"
        elif measurement.provisional:
            state = "provisional"
        else:
            state = "interpolated"
        return track.wrist.copy(), track.wrist + track.hand_shape, state

    track.missed_frames += 1
    if track.wrist is None:
        return None, None, "reachable_long"
    if track.missed_frames <= short_occlusion_frames:
        track.wrist = track.wrist + track.velocity
        track.velocity *= 0.86
        return track.wrist.copy(), track.wrist + track.hand_shape, "predicted_short"
    if track.missed_frames <= medium_occlusion_frames:
        return None, None, "reachable_medium"
    return None, None, "reachable_long"


def temporal_recovery_measurement(
    keypoints: np.ndarray,
    metadata: dict,
    side: str,
    bbox: list[float],
    track: EndpointTrack,
    medium_occlusion_frames: int,
) -> Measurement | None:
    """Recover a weak current-frame hand only when history strongly supports it."""
    if track.wrist is None or track.missed_frames > medium_occlusion_frames:
        return None
    if metadata.get("hand_state") not in {"conflict", "inferred"}:
        return None
    if not metadata.get("selected_hand_valid", False):
        return None
    if float(metadata.get("selected_hand_score", 0.0)) < 2.0:
        return None
    link = float(metadata.get("hand_selection", {}).get("link_to_selected_limb", 1.0))
    if link > 0.08:
        return None

    start, end = HAND_SLICES[side]
    hand = np.asarray(keypoints[start:end], dtype=np.float64)
    if hand.shape != (end - start, 2) or not np.isfinite(hand).all():
        return None
    prediction = track.wrist + track.velocity
    diagonal = bbox_diagonal(bbox)
    gate = min(0.20, 0.05 + 0.008 * max(track.missed_frames, 1)) * diagonal
    if float(np.linalg.norm(hand[0] - prediction)) > gate:
        return None
    return Measurement(
        keypoints=hand,
        precise=False,
        provisional=True,
        source="temporal_recovery",
        quality=float(metadata.get("selected_hand_score", 0.0)),
    )


def hand_span_px(hand: np.ndarray) -> float:
    return math.hypot(float(np.ptp(hand[:, 0])), float(np.ptp(hand[:, 1])))


def bbox_area(bbox: list[float]) -> float:
    return max((bbox[2] - bbox[0]) * (bbox[3] - bbox[1]), 1.0)


def reset_endpoint_track(track: EndpointTrack) -> None:
    track.wrist = None
    track.velocity = None
    track.hand_shape = None
    track.last_observed_frame = None
    track.missed_frames = 0
    track.last_elbow = None
    track.last_direction = None


def resolve_global_hand_assignments(
    frame_measurements: dict[tuple[int, str], Measurement | None],
    tracks: dict[tuple[int, str], EndpointTrack],
    boxes: list[list[float]],
    medium_occlusion_frames: int,
) -> tuple[dict[tuple[int, str], Measurement | None], dict[tuple[int, str], dict]]:
    """Resolve overlapping-person hands before binding them to body sides.

    Large foreground top-down crops may include a smaller background person.
    The same physical hand can therefore be emitted for both people, or only by
    the smaller crop even though it belongs to the foreground person.  Direct
    duplicates are kept on the foreground track when their apparent scales
    agree.  A missing foreground side can then reclaim a nearby candidate using
    its recent hand trajectory and historical hand scale.
    """
    resolved = dict(frame_measurements)
    events: dict[tuple[int, str], dict] = {}
    keys = list(resolved)

    # First collapse direct cross-person duplicates.  Keep this deliberately
    # asymmetric: a substantially larger foreground crop must also carry the
    # substantially larger historical hand shape.
    for index, first_key in enumerate(keys):
        first = resolved[first_key]
        if first is None:
            continue
        for second_key in keys[index + 1 :]:
            second = resolved[second_key]
            if second is None or first_key[0] == second_key[0]:
                continue
            first_hand = first.keypoints
            second_hand = second.keypoints
            first_span = hand_span_px(first_hand)
            second_span = hand_span_px(second_hand)
            center_distance = float(
                np.linalg.norm(np.mean(first_hand, axis=0) - np.mean(second_hand, axis=0))
            )
            if center_distance > max(70.0, 0.5 * max(first_span, second_span)):
                continue

            first_area = bbox_area(boxes[first_key[0]])
            second_area = bbox_area(boxes[second_key[0]])
            if first_area >= 2.5 * second_area and (
                first_span >= 1.45 * second_span
                or first.quality >= second.quality + 0.6
            ):
                owner_key, released_key = first_key, second_key
            elif second_area >= 2.5 * first_area and (
                second_span >= 1.45 * first_span
                or second.quality >= first.quality + 0.6
            ):
                owner_key, released_key = second_key, first_key
            else:
                continue
            resolved[released_key] = None
            reset_endpoint_track(tracks[released_key])
            events[owner_key] = {
                "method": "foreground_duplicate_owner",
                "released_person_id": released_key[0] + 1,
                "released_side": released_key[1],
            }
            events[released_key] = {
                "method": "reassigned_away",
                "assigned_person_id": owner_key[0] + 1,
                "assigned_side": owner_key[1],
            }

    # Then recover missing foreground hands from another crop.  The candidate
    # must fall inside the target's recent motion gate and be much smaller than
    # the target's historical foreground hand shape.
    available_sources = [key for key, value in resolved.items() if value is not None]
    target_candidates = []
    for target_key, target_measurement in resolved.items():
        if target_measurement is not None:
            continue
        target_track = tracks[target_key]
        if (
            target_track.wrist is None
            or target_track.hand_shape is None
            or target_track.missed_frames > medium_occlusion_frames
        ):
            continue
        target_area = bbox_area(boxes[target_key[0]])
        predicted_center = target_track.wrist + np.mean(target_track.hand_shape, axis=0)
        target_span = hand_span_px(target_track.hand_shape)
        gate = 0.20 * bbox_diagonal(boxes[target_key[0]])
        for source_key in available_sources:
            if source_key[0] == target_key[0]:
                continue
            source = resolved[source_key]
            if source is None or source.source != "model":
                continue
            source_area = bbox_area(boxes[source_key[0]])
            if target_area < 2.5 * source_area:
                continue
            source_center = np.mean(source.keypoints, axis=0)
            source_span = hand_span_px(source.keypoints)
            distance = float(np.linalg.norm(source_center - predicted_center))
            if distance > gate or target_span < 1.45 * max(source_span, 1.0):
                continue
            target_candidates.append(
                (distance / gate, target_key, source_key, source_center)
            )

    used_targets = set()
    used_sources = set()
    for _, target_key, source_key, source_center in sorted(target_candidates):
        if target_key in used_targets or source_key in used_sources:
            continue
        target_track = tracks[target_key]
        shape_center = np.mean(target_track.hand_shape, axis=0)
        reassigned_wrist = source_center - shape_center
        reassigned_hand = reassigned_wrist + target_track.hand_shape
        resolved[target_key] = Measurement(
            keypoints=reassigned_hand,
            precise=False,
            provisional=True,
            source="global_reassignment",
            quality=resolved[source_key].quality,
        )
        resolved[source_key] = None
        reset_endpoint_track(tracks[source_key])
        events[target_key] = {
            "method": "foreground_temporal_reassignment",
            "source_person_id": source_key[0] + 1,
            "source_side": source_key[1],
            "observed_center": np.round(source_center, 3).tolist(),
        }
        events[source_key] = {
            "method": "reassigned_away",
            "assigned_person_id": target_key[0] + 1,
            "assigned_side": target_key[1],
        }
        used_targets.add(target_key)
        used_sources.add(source_key)
    return resolved, events


def track_direction(track: EndpointTrack) -> np.ndarray:
    if track.last_direction is not None:
        return track.last_direction / max(float(np.linalg.norm(track.last_direction)), 1e-6)
    if track.velocity is not None and float(np.linalg.norm(track.velocity)) > 1e-6:
        return track.velocity / np.linalg.norm(track.velocity)
    return np.asarray([1.0, 0.0], dtype=np.float64)


def ellipse_polygon(
    center: np.ndarray,
    major_axis: float,
    minor_axis: float,
    direction: np.ndarray,
    point_count: int = 64,
) -> np.ndarray:
    perpendicular = np.asarray([-direction[1], direction[0]], dtype=np.float64)
    angles = np.linspace(0.0, 2.0 * math.pi, point_count, endpoint=False)
    return (
        center
        + np.cos(angles)[:, None] * major_axis * direction
        + np.sin(angles)[:, None] * minor_axis * perpendicular
    )


def clip_polygon_from_anchor_to_circle(
    points: np.ndarray,
    anchor: np.ndarray,
    circle_center: np.ndarray,
    radius: float,
) -> np.ndarray:
    """Clip radial polygon samples to a shoulder-centered anatomical limit."""
    clipped = []
    offset = anchor - circle_center
    for point in points:
        if float(np.linalg.norm(point - circle_center)) <= radius:
            clipped.append(point)
            continue
        ray = point - anchor
        a = float(np.dot(ray, ray))
        b = 2.0 * float(np.dot(offset, ray))
        c = float(np.dot(offset, offset) - radius * radius)
        discriminant = max(b * b - 4.0 * a * c, 0.0)
        roots = [
            value
            for value in (
                (-b - math.sqrt(discriminant)) / max(2.0 * a, 1e-6),
                (-b + math.sqrt(discriminant)) / max(2.0 * a, 1e-6),
            )
            if 0.0 <= value <= 1.0
        ]
        clipped.append(anchor + (max(roots) if roots else 0.0) * ray)
    return np.asarray(clipped, dtype=np.float64)


def medium_motion_region(
    shoulder: np.ndarray,
    upper_length: float,
    forearm_length: float,
    bbox: list[float],
    track: EndpointTrack,
    short_occlusion_frames: int,
) -> dict:
    diagonal = bbox_diagonal(bbox)
    direction = track_direction(track)
    hand_shape = (
        track.hand_shape
        if track.hand_shape is not None
        else np.zeros((21, 2), dtype=np.float64)
    )
    shape_center = np.mean(hand_shape, axis=0)
    centered_shape = hand_shape - shape_center
    perpendicular = np.asarray([-direction[1], direction[0]], dtype=np.float64)
    hand_major = float(np.max(np.abs(centered_shape @ direction)))
    hand_minor = float(np.max(np.abs(centered_shape @ perpendicular)))

    extra_frames = max(track.missed_frames - short_occlusion_frames, 1)
    motion_major = min(0.20, 0.04 + 0.007 * extra_frames) * diagonal
    motion_minor = min(0.13, 0.03 + 0.0045 * extra_frames) * diagonal
    major_axis = max(hand_major + motion_major, 0.08 * diagonal)
    minor_axis = max(hand_minor + motion_minor, 0.06 * diagonal)

    center = track.wrist + shape_center
    tracked_hand_extent = float(np.max(np.linalg.norm(centered_shape, axis=1)))
    hand_reach = max(tracked_hand_extent, 0.10 * diagonal)
    anatomical_reach = upper_length + forearm_length + hand_reach + 0.02 * diagonal
    shoulder_to_center = center - shoulder
    center_distance = float(np.linalg.norm(shoulder_to_center))
    if center_distance > 0.96 * anatomical_reach:
        center = shoulder + shoulder_to_center / center_distance * 0.96 * anatomical_reach

    polygon = ellipse_polygon(center, major_axis, minor_axis, direction)
    polygon = clip_polygon_from_anchor_to_circle(
        polygon,
        center,
        shoulder,
        anatomical_reach,
    )
    direction_deg = math.degrees(math.atan2(float(direction[1]), float(direction[0])))
    return {
        "geometry": "motion_corridor",
        "center": np.round(center, 3).tolist(),
        "predicted_wrist": np.round(track.wrist, 3).tolist(),
        "major_axis_px": round(major_axis, 3),
        "minor_axis_px": round(minor_axis, 3),
        "outer_radius_px": round(major_axis, 3),
        "direction_deg": round(direction_deg, 3),
        "max_reach_from_shoulder_px": round(anatomical_reach, 3),
        "polygon": np.round(polygon, 3).tolist(),
        "requires_door_intersection": True,
        "render_full_region": True,
        "distance_lower_bound_rule": "minimum distance from motion-corridor polygon to door_edge",
    }


def long_reachable_region(
    shoulder: np.ndarray,
    upper_length: float,
    forearm_length: float,
    bbox: list[float],
    track: EndpointTrack,
    medium_occlusion_frames: int,
) -> dict:
    diagonal = bbox_diagonal(bbox)
    height = max(bbox[3] - bbox[1], 1.0)
    arm_reach = upper_length + forearm_length
    hand_reach = min(0.13 * diagonal, 0.17 * height)
    safety_margin = 0.025 * diagonal
    outer_radius = arm_reach + hand_reach + safety_margin
    inner_radius = max(abs(upper_length - forearm_length) - 0.04 * diagonal, 0.0)
    if track.last_direction is None:
        direction_deg = 0.0
        half_angle_deg = 180.0
    else:
        direction = track_direction(track)
        direction_deg = math.degrees(math.atan2(float(direction[1]), float(direction[0])))
        half_angle_deg = min(
            180.0,
            65.0 + 2.0 * max(track.missed_frames - medium_occlusion_frames, 0),
        )
    return {
        "geometry": "shoulder_sector",
        "center": np.round(shoulder, 3).tolist(),
        "inner_radius_px": round(inner_radius, 3),
        "outer_radius_px": round(outer_radius, 3),
        "arm_reach_px": round(arm_reach, 3),
        "hand_reach_px": round(hand_reach, 3),
        "safety_margin_px": round(safety_margin, 3),
        "direction_deg": round(direction_deg, 3),
        "half_angle_deg": round(half_angle_deg, 3),
        "requires_door_intersection": True,
        "render_full_region": False,
        "distance_lower_bound_rule": "intersect shoulder sector with door danger band before alarming",
    }


def reachable_region(
    shoulder: np.ndarray,
    upper_length: float,
    forearm_length: float,
    bbox: list[float],
    track: EndpointTrack,
    state: str = "reachable_long",
    short_occlusion_frames: int = 12,
    medium_occlusion_frames: int = 30,
) -> dict:
    if state == "reachable_medium" and track.wrist is not None:
        return medium_motion_region(
            shoulder,
            upper_length,
            forearm_length,
            bbox,
            track,
            short_occlusion_frames,
        )
    return long_reachable_region(
        shoulder,
        upper_length,
        forearm_length,
        bbox,
        track,
        medium_occlusion_frames,
    )


def calibrated_bone_lengths(
    upper_ratio: float,
    forearm_ratio: float,
    bbox: list[float],
) -> tuple[float, float]:
    diagonal = bbox_diagonal(bbox)
    height = max(bbox[3] - bbox[1], 1.0)
    upper = upper_ratio * diagonal
    forearm = forearm_ratio * diagonal
    total = upper + forearm
    # Keep only a loose sanity cap.  The previous 0.55 * bbox-height cap was
    # unsafe for close, clipped CMS views: a valid projected arm can occupy
    # almost the full detection height.  The ratios are already robust medians,
    # so this cap only limits catastrophic calibration outliers.
    maximum_total = 0.95 * diagonal
    if total > maximum_total:
        scale = maximum_total / total
        upper *= scale
        forearm *= scale
    return upper, forearm