"""Runtime pose fusion helpers."""
from __future__ import annotations
import math
from dataclasses import dataclass, field
import cv2
import numpy as np
from pose_geometry import BODY_WRISTS, HAND_EDGES_LOCAL, HAND_SLICES, normalized_keypoints


BODY_SIDE_INDICES = {"left": (5, 7, 9), "right": (6, 8, 10)}


SIDE_COLORS = {"left": (40, 40, 255), "right": (30, 230, 30)}


HAND_COLORS = {"left": (255, 80, 220), "right": (255, 190, 40)}


@dataclass
class Candidate:
    source: str
    keypoints: np.ndarray
    scores: np.ndarray
    normalized: np.ndarray
    limb_score: float
    hand_score: float
    upper_length: float
    forearm_length: float
    hand_span: float
    wrist_link: float
    limb_valid: bool
    hand_valid: bool


@dataclass
class SideTrack:
    limb_previous: np.ndarray | None = None
    limb_velocity: np.ndarray | None = None
    hand_previous: np.ndarray | None = None
    hand_velocity: np.ndarray | None = None
    limb_source: str | None = None
    hand_source: str | None = None
    previous_hand_eligible: bool = False


@dataclass
class PersonTrack:
    sides: dict[str, SideTrack] = field(
        default_factory=lambda: {"left": SideTrack(), "right": SideTrack()}
    )
    smoothed: 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 build_candidate(
    source: str,
    keypoints: np.ndarray,
    scores: np.ndarray,
    bbox: list[float],
    side: str,
) -> Candidate:
    normalized = normalized_keypoints(keypoints, bbox)
    shoulder, elbow, wrist = BODY_SIDE_INDICES[side]
    hand_start, hand_end = HAND_SLICES[side]
    hand = keypoints[hand_start:hand_end]
    diagonal = bbox_diagonal(bbox)
    upper_length = float(np.linalg.norm(normalized[shoulder] - normalized[elbow]))
    forearm_length = float(np.linalg.norm(normalized[elbow] - normalized[wrist]))
    hand_span = math.hypot(float(np.ptp(hand[:, 0])), float(np.ptp(hand[:, 1]))) / diagonal
    wrist_link = float(np.linalg.norm(keypoints[wrist] - hand[0])) / diagonal
    limb_score = float(np.mean(scores[[shoulder, elbow, wrist]]))
    hand_score = float(np.mean(scores[hand_start:hand_end]))
    limb_valid = 0.025 <= upper_length <= 0.8 and 0.025 <= forearm_length <= 0.8
    hand_valid = 0.008 <= hand_span <= 0.35 and wrist_link <= 0.12
    return Candidate(
        source=source,
        keypoints=keypoints,
        scores=scores,
        normalized=normalized,
        limb_score=limb_score,
        hand_score=hand_score,
        upper_length=upper_length,
        forearm_length=forearm_length,
        hand_span=hand_span,
        wrist_link=wrist_link,
        limb_valid=limb_valid,
        hand_valid=hand_valid,
    )


def predicted(previous: np.ndarray | None, velocity: np.ndarray | None) -> np.ndarray | None:
    if previous is None:
        return None
    if velocity is None:
        return previous
    return previous + 0.65 * velocity


def robust_distance(points: np.ndarray, target: np.ndarray | None) -> float:
    if target is None:
        return 0.0
    distances = np.linalg.norm(points - target, axis=1)
    return float(np.median(distances))


def select_limb(
    candidates: list[Candidate],
    side: str,
    track: SideTrack | None,
) -> tuple[Candidate, dict[str, float]]:
    indices = np.asarray(BODY_SIDE_INDICES[side])
    target = predicted(track.limb_previous, track.limb_velocity) if track else None
    scored: list[tuple[float, Candidate, float]] = []
    for candidate in candidates:
        temporal_cost = robust_distance(candidate.normalized[indices], target)
        geometry_bonus = 1.2 if candidate.limb_valid else -5.0
        switch_bonus = 0.35 if track and candidate.source == track.limb_source else 0.0
        score = candidate.limb_score + geometry_bonus + switch_bonus - 7.0 * temporal_cost
        scored.append((score, candidate, temporal_cost))
    score, selected, temporal_cost = max(scored, key=lambda item: item[0])
    return selected, {"score": score, "temporal_cost": temporal_cost}


def hand_link_to_limb(
    candidate: Candidate,
    selected_limb: Candidate,
    bbox: list[float],
    side: str,
) -> float:
    hand_start, _ = HAND_SLICES[side]
    wrist_idx = BODY_WRISTS[side]
    return float(
        np.linalg.norm(candidate.keypoints[hand_start] - selected_limb.keypoints[wrist_idx])
    ) / bbox_diagonal(bbox)


def select_hand(
    candidates: list[Candidate],
    selected_limb: Candidate,
    bbox: list[float],
    side: str,
    track: SideTrack | None,
) -> tuple[Candidate, bool, dict[str, float]]:
    start, end = HAND_SLICES[side]
    target = predicted(track.hand_previous, track.hand_velocity) if track else None
    scored: list[tuple[float, Candidate, bool, float, float]] = []
    for candidate in candidates:
        temporal_cost = robust_distance(candidate.normalized[start:end], target)
        link = hand_link_to_limb(candidate, selected_limb, bbox, side)
        valid = candidate.hand_valid
        geometry_bonus = 3.0 if valid else -7.0
        switch_bonus = 0.45 if track and candidate.source == track.hand_source else 0.0
        limb_source_bonus = 0.15 if candidate.source == selected_limb.source else 0.0
        # A newly visible endpoint may disagree with the currently selected
        # (occlusion-driven) limb.  Keep this as a soft cost so strong image
        # evidence can switch both the hand and its limb together.
        link_penalty = 1.5 * min(max(link - 0.12, 0.0), 0.2)
        score = (
            candidate.hand_score
            + geometry_bonus
            + switch_bonus
            + limb_source_bonus
            - link_penalty
            - 8.0 * min(temporal_cost, 0.15)
        )
        scored.append((score, candidate, valid, temporal_cost, link))
    score, selected, valid, temporal_cost, link = max(scored, key=lambda item: item[0])
    return selected, valid, {
        "score": score,
        "temporal_cost": temporal_cost,
        "link_to_selected_limb": link,
    }


def candidate_disagreement(
    raw: Candidate,
    isolated: Candidate,
    side: str,
) -> tuple[float, float]:
    shoulder, elbow, wrist = BODY_SIDE_INDICES[side]
    start, end = HAND_SLICES[side]
    limb = float(
        np.median(
            np.linalg.norm(
                raw.normalized[[shoulder, elbow, wrist]]
                - isolated.normalized[[shoulder, elbow, wrist]],
                axis=1,
            )
        )
    )
    hand_distances = np.linalg.norm(
        raw.normalized[start:end] - isolated.normalized[start:end], axis=1
    )
    hand = float(np.percentile(hand_distances, 90))
    return limb, hand


def suppress_duplicate_hands(
    fused: np.ndarray,
    metadata: dict,
    bbox: list[float],
) -> None:
    left_start, left_end = HAND_SLICES["left"]
    right_start, right_end = HAND_SLICES["right"]
    left_center = np.mean(fused[left_start:left_end], axis=0)
    right_center = np.mean(fused[right_start:right_end], axis=0)
    separation = float(np.linalg.norm(left_center - right_center)) / bbox_diagonal(bbox)
    if separation >= 0.08:
        return
    weaker = min(
        ("left", "right"),
        key=lambda side: metadata[side]["selected_hand_score"],
    )
    metadata[weaker]["hand_state"] = "duplicate"
    metadata[weaker]["distance_eligible"] = False
    metadata[weaker]["alarm_eligible"] = False
    metadata[weaker]["duplicate_hand_separation"] = round(separation, 5)


def fuse_person(
    raw_keypoints: np.ndarray,
    raw_scores: np.ndarray,
    isolated_keypoints: np.ndarray,
    isolated_scores: np.ndarray,
    bbox: list[float],
    track: PersonTrack | None,
) -> tuple[np.ndarray, dict]:
    fused = raw_keypoints.copy()
    metadata: dict[str, dict] = {}

    for side in ("left", "right"):
        side_track = track.sides[side] if track else None
        raw = build_candidate("raw", raw_keypoints, raw_scores, bbox, side)
        isolated = build_candidate(
            "isolated", isolated_keypoints, isolated_scores, bbox, side
        )
        candidates = [raw, isolated]
        limb, limb_detail = select_limb(candidates, side, side_track)
        hand, hand_valid, hand_detail = select_hand(
            candidates, limb, bbox, side, side_track
        )

        other_hand = isolated if hand.source == "raw" else raw
        endpoint_advantage = hand.hand_score - other_hand.hand_score
        if hand.source != limb.source and hand_valid and endpoint_advantage >= 0.6:
            limb = hand
            limb_detail["visible_endpoint_override"] = 1.0
            hand_detail["link_to_selected_limb"] = hand.wrist_link

        body_indices = np.asarray(BODY_SIDE_INDICES[side])
        hand_start, hand_end = HAND_SLICES[side]
        fused[body_indices] = limb.keypoints[body_indices]
        fused[hand_start:hand_end] = hand.keypoints[hand_start:hand_end]

        limb_disagreement, hand_disagreement = candidate_disagreement(raw, isolated, side)
        if limb_disagreement <= 0.05:
            limb_state = "consensus"
        else:
            limb_state = "inferred"

        score_advantage = raw.hand_score - isolated.hand_score
        raw_temporally_stable = hand_detail["temporal_cost"] <= 0.08
        if not hand_valid:
            hand_state = "rejected"
            distance_eligible = False
        elif hand_disagreement <= 0.06:
            hand_state = "consensus"
            distance_eligible = True
        elif (
            hand.source == "raw"
            and score_advantage >= 0.6
            and raw_temporally_stable
        ):
            hand_state = "visible_endpoint"
            distance_eligible = True
        elif hand.source == "isolated":
            hand_state = "inferred"
            distance_eligible = False
        else:
            hand_state = "conflict"
            distance_eligible = False
        fast_alarm_endpoint = (
            hand_valid and hand.source == "raw" and score_advantage >= 0.6
        )
        alarm_eligible = distance_eligible or fast_alarm_endpoint

        metadata[side] = {
            "limb_source": limb.source,
            "hand_source": hand.source,
            "limb_state": limb_state,
            "hand_state": hand_state,
            "distance_eligible": distance_eligible,
            "alarm_eligible": alarm_eligible,
            "limb_disagreement": round(limb_disagreement, 5),
            "hand_disagreement": round(hand_disagreement, 5),
            "raw_hand_score": round(raw.hand_score, 4),
            "isolated_hand_score": round(isolated.hand_score, 4),
            "selected_hand_span": round(hand.hand_span, 5),
            "selected_hand_score": round(hand.hand_score, 4),
            "selected_hand_valid": hand_valid,
            "limb_selection": {key: round(value, 5) for key, value in limb_detail.items()},
            "hand_selection": {key: round(value, 5) for key, value in hand_detail.items()},
        }

        if side_track is not None:
            limb_now = limb.normalized[body_indices]
            hand_now = hand.normalized[hand_start:hand_end]
            if side_track.limb_previous is not None:
                delta = limb_now - side_track.limb_previous
                side_track.limb_velocity = (
                    delta
                    if side_track.limb_velocity is None
                    else 0.65 * side_track.limb_velocity + 0.35 * delta
                )
            if hand_valid:
                if side_track.hand_previous is not None:
                    delta = hand_now - side_track.hand_previous
                    side_track.hand_velocity = (
                        delta
                        if side_track.hand_velocity is None
                        else 0.65 * side_track.hand_velocity + 0.35 * delta
                    )
                side_track.hand_previous = hand_now
                side_track.hand_source = hand.source
            side_track.limb_previous = limb_now
            side_track.limb_source = limb.source
            side_track.previous_hand_eligible = distance_eligible

    suppress_duplicate_hands(fused, metadata, bbox)
    return fused, metadata


def adaptive_smooth(
    keypoints: np.ndarray,
    bbox: list[float],
    track: PersonTrack,
) -> np.ndarray:
    normalized = normalized_keypoints(keypoints, bbox)
    if track.smoothed is None:
        track.smoothed = normalized
    else:
        displacement = np.linalg.norm(normalized - track.smoothed, axis=1)
        alpha = 0.32 + 0.58 * np.clip(displacement / 0.08, 0.0, 1.0)
        track.smoothed = alpha[:, None] * normalized + (1.0 - alpha[:, None]) * track.smoothed

    x1, y1, x2, y2 = bbox
    result = track.smoothed.copy()
    result[:, 0] = result[:, 0] * max(x2 - x1, 1.0) + x1
    result[:, 1] = result[:, 1] * max(y2 - y1, 1.0) + y1
    return result


def build_global_hands(
    keypoints: list[np.ndarray],
    metadata: list[dict],
) -> list[dict]:
    candidates = []
    for person_idx, (person_keypoints, person_meta) in enumerate(
        zip(keypoints, metadata), start=1
    ):
        for side, (start, end) in HAND_SLICES.items():
            meta = person_meta[side]
            if not meta["alarm_eligible"]:
                continue
            hand = person_keypoints[start:end]
            span_px = math.hypot(float(np.ptp(hand[:, 0])), float(np.ptp(hand[:, 1])))
            candidates.append(
                {
                    "person_id": person_idx,
                    "side": side,
                    "source": meta["hand_source"],
                    "state": meta["hand_state"],
                    "distance_eligible": meta["distance_eligible"],
                    "alarm_eligible": meta["alarm_eligible"],
                    "score": meta["selected_hand_score"],
                    "center": np.mean(hand, axis=0),
                    "span_px": span_px,
                    "keypoints": hand,
                    "aliases": [],
                }
            )

    kept: list[dict] = []
    for candidate in sorted(candidates, key=lambda item: item["score"], reverse=True):
        duplicate = None
        for existing in kept:
            threshold = float(np.clip(0.45 * max(candidate["span_px"], existing["span_px"]), 18, 70))
            if float(np.linalg.norm(candidate["center"] - existing["center"])) <= threshold:
                duplicate = existing
                break
        if duplicate is None:
            kept.append(candidate)
        else:
            duplicate["aliases"].append(
                {"person_id": candidate["person_id"], "side": candidate["side"]}
            )

    return [
        {
            "person_id": item["person_id"],
            "side": item["side"],
            "source": item["source"],
            "state": item["state"],
            "distance_eligible": item["distance_eligible"],
            "alarm_eligible": item["alarm_eligible"],
            "score": item["score"],
            "center": np.round(item["center"], 3).tolist(),
            "span_px": round(item["span_px"], 3),
            "keypoints": np.round(item["keypoints"], 3).tolist(),
            "aliases": item["aliases"],
        }
        for item in kept
    ]


def draw_dashed_line(
    image: np.ndarray,
    start: np.ndarray,
    end: np.ndarray,
    color: tuple[int, int, int],
    thickness: int = 2,
) -> None:
    a = np.asarray(start, dtype=np.float64)
    b = np.asarray(end, dtype=np.float64)
    length = float(np.linalg.norm(b - a))
    if length < 1.0:
        return
    direction = (b - a) / length
    for offset in np.arange(0.0, length, 14.0):
        segment_start = a + direction * offset
        segment_end = a + direction * min(offset + 8.0, length)
        cv2.line(
            image,
            tuple(np.round(segment_start).astype(int)),
            tuple(np.round(segment_end).astype(int)),
            color,
            thickness,
            cv2.LINE_AA,
        )


def draw_fused(
    image: np.ndarray,
    boxes: list[list[float]],
    keypoints: list[np.ndarray],
    metadata: list[dict],
    title: str,
) -> np.ndarray:
    canvas = image.copy()
    for person_idx, (bbox, kpts, person_meta) in enumerate(
        zip(boxes, keypoints, metadata), start=1
    ):
        x1, y1, x2, y2 = (int(round(value)) for value in bbox)
        cv2.rectangle(canvas, (x1, y1), (x2, y2), (0, 220, 255), 2)
        cv2.putText(
            canvas,
            f"P{person_idx}",
            (x1 + 4, max(22, y1 + 24)),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.72,
            (0, 220, 255),
            2,
            cv2.LINE_AA,
        )
        for side in ("left", "right"):
            shoulder, elbow, wrist = BODY_SIDE_INDICES[side]
            side_meta = person_meta[side]
            color = SIDE_COLORS[side]
            for a, b in ((shoulder, elbow), (elbow, wrist)):
                if side_meta["limb_state"] == "consensus":
                    cv2.line(
                        canvas,
                        tuple(np.round(kpts[a]).astype(int)),
                        tuple(np.round(kpts[b]).astype(int)),
                        color,
                        3,
                        cv2.LINE_AA,
                    )
                else:
                    draw_dashed_line(canvas, kpts[a], kpts[b], (0, 140, 255), 3)

            start, end = HAND_SLICES[side]
            hand = kpts[start:end]
            if side_meta["hand_state"] != "rejected":
                hand_color = HAND_COLORS[side]
                if not side_meta["distance_eligible"]:
                    hand_color = (0, 140, 255)
                for a, b in HAND_EDGES_LOCAL:
                    if side_meta["distance_eligible"]:
                        cv2.line(
                            canvas,
                            tuple(np.round(hand[a]).astype(int)),
                            tuple(np.round(hand[b]).astype(int)),
                            hand_color,
                            2,
                            cv2.LINE_AA,
                        )
                    else:
                        draw_dashed_line(canvas, hand[a], hand[b], hand_color, 2)
            label = (
                f"{side[0].upper()}:{side_meta['hand_state'][:4]}"
                f"/{side_meta['hand_source'][0]}"
            )
            label_y = max(45, y1 + (48 if side == "left" else 70))
            cv2.putText(
                canvas,
                label,
                (x1 + 4, label_y),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                HAND_COLORS[side] if side_meta["distance_eligible"] else (0, 140, 255),
                1,
                cv2.LINE_AA,
            )

    cv2.rectangle(canvas, (0, 0), (650, 45), (0, 0, 0), -1)
    cv2.putText(
        canvas,
        title,
        (10, 31),
        cv2.FONT_HERSHEY_SIMPLEX,
        0.8,
        (255, 255, 255),
        2,
        cv2.LINE_AA,
    )
    return canvas