from __future__ import annotations

import copy
import json
import math
from pathlib import Path
from typing import Any, Iterable, Mapping, Sequence


def clamp(value: float, low: float, high: float) -> float:
    return max(low, min(high, float(value)))


def lerp(a: float, b: float, ratio: float) -> float:
    return (1.0 - ratio) * float(a) + ratio * float(b)


def vec_sub(a: tuple[float, float], b: tuple[float, float]) -> tuple[float, float]:
    return (a[0] - b[0], a[1] - b[1])


def vec_mid(a: tuple[float, float], b: tuple[float, float]) -> tuple[float, float]:
    return ((a[0] + b[0]) * 0.5, (a[1] + b[1]) * 0.5)


def vec_len(a: tuple[float, float]) -> float:
    return math.hypot(a[0], a[1])


def vec_norm(a: tuple[float, float], fallback: tuple[float, float]) -> tuple[float, float]:
    length = vec_len(a)
    if length <= 1e-6:
        return fallback
    return (a[0] / length, a[1] / length)


def vec_dot(a: tuple[float, float], b: tuple[float, float]) -> float:
    return a[0] * b[0] + a[1] * b[1]


def angle_between(a: tuple[float, float], b: tuple[float, float], fallback_axis: tuple[float, float]) -> float:
    na = vec_norm(a, fallback_axis)
    nb = vec_norm(b, fallback_axis)
    return math.acos(clamp(vec_dot(na, nb), -1.0, 1.0))


def load_home_pose(path: str | Path, expected_joints: int) -> list[float]:
    last_q: list[float] | None = None
    with open(path, "r", encoding="utf-8") as handle:
        for line in handle:
            try:
                data = json.loads(line)
            except json.JSONDecodeError:
                continue
            q = data.get("q")
            if isinstance(q, list) and len(q) == expected_joints:
                last_q = [float(v) for v in q]

    if last_q is None:
        raise ValueError(f"No valid {expected_joints}-joint pose found in {path}")
    return last_q


class PoseToG1Mapper:
    """
    Converts 2D YOLO pose keypoints into a replay-compatible G1 joint frame.
    """

    def __init__(
        self,
        home_q: Sequence[float],
        joint_config: Mapping[str, Any],
        pose_mapping_config: Mapping[str, Any],
        *,
        min_conf: float | None = None,
        smoothing: float | None = None,
    ) -> None:
        self.joint_config = dict(joint_config)
        self.pose_mapping_config = dict(pose_mapping_config)
        self.num_motors = int(self.joint_config["g1_num_motor"])
        if len(home_q) != self.num_motors:
            raise ValueError(f"Expected {self.num_motors} joints in home_q, got {len(home_q)}")

        self.home_q = [float(v) for v in home_q]
        self.prev_q = list(self.home_q)

        waist_cfg = self.joint_config["waist_joint_indices"]
        self.waist_yaw = int(waist_cfg["yaw"])
        self.waist_roll = int(waist_cfg["roll"])
        self.waist_pitch = int(waist_cfg["pitch"])

        self.joint_limits = {
            int(idx): (float(bounds[0]), float(bounds[1]))
            for idx, bounds in self.joint_config["joint_limits"].items()
        }
        self.keypoint_ids = {
            name: int(value)
            for name, value in self.pose_mapping_config["body_keypoint_indices"].items()
        }
        self.serializable_indices = [int(idx) for idx in self.pose_mapping_config["serializable_keypoint_indices"]]
        fallback_cfg = self.pose_mapping_config["fallback_vectors"]
        self.axis_x_fallback = tuple(float(v) for v in fallback_cfg["axis_x"])
        self.torso_up_fallback = tuple(float(v) for v in fallback_cfg["torso_up"])
        self.diagnostics_defaults = copy.deepcopy(self.pose_mapping_config["diagnostics_defaults"])
        self.min_segment_length = float(self.pose_mapping_config["min_segment_length_px"])
        self.waist_map = self.pose_mapping_config["waist"]
        self.arm_maps = self.pose_mapping_config["arms"]

        default_min_conf = float(self.pose_mapping_config["min_conf"])
        default_smoothing = float(self.pose_mapping_config["default_smoothing"])
        self.min_conf = float(default_min_conf if min_conf is None else min_conf)
        self.smoothing = clamp(float(default_smoothing if smoothing is None else smoothing), 0.0, 1.0)

    def reset(self) -> None:
        self.prev_q = list(self.home_q)

    def convert(
        self,
        keypoints: Mapping[int, Sequence[float]] | None,
    ) -> tuple[list[float], dict[str, float | int | bool | None]]:
        target_q = list(self.home_q)
        tracked_joint_indices: set[int] = set()

        diagnostics: dict[str, float | int | bool | None] = copy.deepcopy(self.diagnostics_defaults)

        if not keypoints:
            self.prev_q = list(self.home_q)
            return list(self.home_q), diagnostics

        visible = {idx: self._point(keypoints.get(idx)) for idx in self.serializable_indices}
        diagnostics["visible_keypoints"] = sum(1 for value in visible.values() if value is not None)

        left_shoulder = visible.get(self.keypoint_ids["left_shoulder"])
        right_shoulder = visible.get(self.keypoint_ids["right_shoulder"])
        left_hip = visible.get(self.keypoint_ids["left_hip"])
        right_hip = visible.get(self.keypoint_ids["right_hip"])

        if left_shoulder and right_shoulder:
            shoulder_center = vec_mid(left_shoulder, right_shoulder)
            shoulder_axis = vec_norm(vec_sub(right_shoulder, left_shoulder), self.axis_x_fallback)

            if left_hip and right_hip:
                hip_center = vec_mid(left_hip, right_hip)
                torso_up = vec_norm(vec_sub(shoulder_center, hip_center), self.torso_up_fallback)
                center_offset = clamp(
                    (shoulder_center[0] - hip_center[0]) / max(vec_len(vec_sub(right_shoulder, left_shoulder)), 1.0),
                    -1.0,
                    1.0,
                )
                yaw_gain = float(self.waist_map["yaw_center_gain"])
                target_q[self.waist_yaw] = self._joint(self.waist_yaw, self.home_q[self.waist_yaw] + yaw_gain * center_offset)
                tracked_joint_indices.add(self.waist_yaw)
            else:
                torso_up = self.torso_up_fallback

            diagnostics["tracked"] = True
            diagnostics["torso_ready"] = True
            roll_gain = float(self.waist_map["roll_shoulder_gain"])
            target_q[self.waist_roll] = self._joint(self.waist_roll, self.home_q[self.waist_roll] + roll_gain * shoulder_axis[1])
            tracked_joint_indices.add(self.waist_roll)

            left_q, left_info = self._arm_from_pose(
                side="left",
                shoulder=left_shoulder,
                elbow=visible.get(self.keypoint_ids["left_elbow"]),
                wrist=visible.get(self.keypoint_ids["left_wrist"]),
                torso_up=torso_up,
                shoulder_axis=shoulder_axis,
            )
            if left_q is not None:
                for idx, value in left_q.items():
                    target_q[idx] = value
                    tracked_joint_indices.add(idx)
                diagnostics["left_ready"] = True
                diagnostics["left_lift"] = left_info["lift"]
                diagnostics["left_bend"] = left_info["bend"]

            right_q, right_info = self._arm_from_pose(
                side="right",
                shoulder=right_shoulder,
                elbow=visible.get(self.keypoint_ids["right_elbow"]),
                wrist=visible.get(self.keypoint_ids["right_wrist"]),
                torso_up=torso_up,
                shoulder_axis=shoulder_axis,
            )
            if right_q is not None:
                for idx, value in right_q.items():
                    target_q[idx] = value
                    tracked_joint_indices.add(idx)
                diagnostics["right_ready"] = True
                diagnostics["right_lift"] = right_info["lift"]
                diagnostics["right_bend"] = right_info["bend"]

        smoothed = []
        for idx, (prev, new) in enumerate(zip(self.prev_q, target_q)):
            if idx not in tracked_joint_indices:
                smoothed.append(self.home_q[idx])
                continue
            blended = lerp(prev, new, self.smoothing)
            if idx in self.joint_limits:
                blended = self._joint(idx, blended)
            smoothed.append(blended)

        self.prev_q = smoothed
        return list(smoothed), diagnostics

    def _point(self, value: Sequence[float] | None) -> tuple[float, float] | None:
        if value is None or len(value) < 3:
            return None
        x, y, conf = float(value[0]), float(value[1]), float(value[2])
        if conf < self.min_conf:
            return None
        return (x, -y)

    def _joint(self, idx: int, value: float) -> float:
        low, high = self.joint_limits[idx]
        return clamp(value, low, high)

    def _arm_from_pose(
        self,
        *,
        side: str,
        shoulder: tuple[float, float] | None,
        elbow: tuple[float, float] | None,
        wrist: tuple[float, float] | None,
        torso_up: tuple[float, float],
        shoulder_axis: tuple[float, float],
    ) -> tuple[dict[int, float] | None, dict[str, float]]:
        if not shoulder or not elbow or not wrist:
            return None, {"lift": 0.0, "bend": 0.0}

        upper = vec_sub(elbow, shoulder)
        fore = vec_sub(wrist, elbow)
        hand = vec_sub(wrist, shoulder)
        if vec_len(upper) < self.min_segment_length or vec_len(fore) < self.min_segment_length:
            return None, {"lift": 0.0, "bend": 0.0}

        arm_cfg = self.arm_maps[side]
        joint_indices = {name: int(idx) for name, idx in arm_cfg["joint_indices"].items()}

        upper_dir = vec_norm(upper, torso_up)
        fore_dir = vec_norm(fore, upper_dir)
        hand_dir = vec_norm(hand, upper_dir)

        side_sign = -1.0 if side == "left" else 1.0
        lift = clamp(vec_dot(upper_dir, torso_up), -1.0, 1.0)
        outward = clamp(side_sign * vec_dot(upper_dir, shoulder_axis), -1.0, 1.0)
        cross_body = clamp(-side_sign * vec_dot(hand_dir, shoulder_axis), -1.0, 1.0)
        fore_lift = clamp(vec_dot(fore_dir, torso_up), -1.0, 1.0)
        bend_ratio = angle_between(upper_dir, fore_dir, self.axis_x_fallback) / math.pi
        horizontal = 1.0 - abs(lift)

        shoulder_pitch_cfg = arm_cfg["shoulder_pitch"]
        shoulder_roll_cfg = arm_cfg["shoulder_roll"]
        shoulder_yaw_cfg = arm_cfg["shoulder_yaw"]
        elbow_cfg = arm_cfg["elbow"]
        wrist_roll_cfg = arm_cfg["wrist_roll"]
        wrist_pitch_cfg = arm_cfg["wrist_pitch"]
        wrist_yaw_cfg = arm_cfg["wrist_yaw"]

        q = {
            joint_indices["shoulder_pitch"]: self._joint(
                joint_indices["shoulder_pitch"],
                lerp(float(shoulder_pitch_cfg["down"]), float(shoulder_pitch_cfg["up"]), (lift + 1.0) * 0.5),
            ),
            joint_indices["shoulder_roll"]: self._joint(
                joint_indices["shoulder_roll"],
                self.home_q[joint_indices["shoulder_roll"]]
                + float(shoulder_roll_cfg["outward_horizontal"]) * outward * horizontal
                + float(shoulder_roll_cfg["lift_positive"]) * max(lift, 0.0)
                + float(shoulder_roll_cfg["cross_body"]) * cross_body,
            ),
            joint_indices["shoulder_yaw"]: self._joint(
                joint_indices["shoulder_yaw"],
                self.home_q[joint_indices["shoulder_yaw"]]
                + float(shoulder_yaw_cfg["cross_body"]) * cross_body
                + float(shoulder_yaw_cfg["fore_lift"]) * fore_lift,
            ),
            joint_indices["elbow"]: self._joint(
                joint_indices["elbow"],
                float(elbow_cfg["base"]) + float(elbow_cfg["bend"]) * bend_ratio,
            ),
            joint_indices["wrist_roll"]: self._joint(
                joint_indices["wrist_roll"],
                self.home_q[joint_indices["wrist_roll"]]
                + float(wrist_roll_cfg["outward"]) * outward
                + float(wrist_roll_cfg["cross_body"]) * cross_body,
            ),
            joint_indices["wrist_pitch"]: self._joint(
                joint_indices["wrist_pitch"],
                self.home_q[joint_indices["wrist_pitch"]] + float(wrist_pitch_cfg["fore_lift"]) * fore_lift,
            ),
            joint_indices["wrist_yaw"]: self._joint(
                joint_indices["wrist_yaw"],
                self.home_q[joint_indices["wrist_yaw"]] + float(wrist_yaw_cfg["cross_body"]) * cross_body,
            ),
        }

        return q, {"lift": lift, "bend": bend_ratio}


def keypoints_to_serializable(
    keypoints: Mapping[int, Sequence[float]] | None,
    indices: Iterable[int],
) -> dict[str, list[float] | None]:
    serializable: dict[str, list[float] | None] = {}
    if not keypoints:
        for idx in indices:
            serializable[str(idx)] = None
        return serializable

    for idx in indices:
        point = keypoints.get(idx)
        if point is None:
            serializable[str(idx)] = None
            continue
        serializable[str(idx)] = [float(point[0]), float(point[1]), float(point[2])]
    return serializable
