from __future__ import annotations

import math
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np

HOST_ROOT = Path(__file__).resolve().parents[1]
PROJECT_ROOT = HOST_ROOT.parents[1]
RL_ROOT = PROJECT_ROOT / 'guguji_rl'
if str(RL_ROOT) not in sys.path:
    sys.path.insert(0, str(RL_ROOT))

from guguji_rl.config import load_config, resolve_project_path
from guguji_rl.urdf_utils import parse_joint_limits


@dataclass(slots=True)
class DeploymentObservation:
    observation: np.ndarray
    joint_targets: np.ndarray


class DeployActionMapper:
    """复用仿真环境里的动作映射逻辑，减少 sim-to-real 的隐藏差异。"""

    def __init__(
        self,
        *,
        rl_config_path: str | Path,
        deploy_config: dict[str, Any],
    ) -> None:
        self.config = load_config(rl_config_path)
        self.robot_config = self.config['robot']
        self.sim_config = self.config['sim']
        self.task_config = self.config['task']
        self.reference_gait_config = self.robot_config.get('reference_gait', {})
        self.reference_gait_enabled = bool(self.reference_gait_config.get('enabled', False))
        self.reference_gait_period = max(
            float(self.reference_gait_config.get('period', 0.9)),
            float(self.sim_config['control_dt']),
        )

        urdf_path = resolve_project_path(self.config, self.robot_config['urdf_path'])
        self.joint_limits = parse_joint_limits(urdf_path, self.robot_config['joint_names'])
        self.joint_names = [joint_limit.name for joint_limit in self.joint_limits]
        self.joint_lower = np.array([joint.lower for joint in self.joint_limits], dtype=np.float32)
        self.joint_upper = np.array([joint.upper for joint in self.joint_limits], dtype=np.float32)
        self.joint_mid = np.array([joint.midpoint for joint in self.joint_limits], dtype=np.float32)
        self.joint_half_range = np.array([joint.half_range for joint in self.joint_limits], dtype=np.float32)

        nominal_joint_config = self.robot_config.get('nominal_joint_positions', {})
        self.nominal_joint_targets = np.array(
            [float(nominal_joint_config.get(joint_name, 0.0)) for joint_name in self.joint_names],
            dtype=np.float32,
        )
        self.action_scale = float(self.robot_config.get('action_scale', 1.0))
        self.action_smoothing = float(np.clip(self.robot_config.get('action_smoothing', 0.0), 0.0, 0.99))
        self.current_action_residual = np.zeros(len(self.joint_names), dtype=np.float32)
        self.previous_action = np.zeros(len(self.joint_names), dtype=np.float32)
        self.gait_phase = 0.0

        self.estimated_base_height = float(deploy_config.get('estimated_base_height_m', 0.35))
        self.estimated_forward_velocity = float(deploy_config.get('estimated_forward_velocity_m_s', 0.0))
        self.target_forward_velocity = float(
            deploy_config.get(
                'target_forward_velocity_m_s',
                self.task_config.get('target_forward_velocity', 0.0),
            )
        )

    @property
    def control_dt(self) -> float:
        return float(self.sim_config['control_dt'])

    def reset(self) -> None:
        self.current_action_residual[:] = 0.0
        self.previous_action[:] = 0.0
        self.gait_phase = 0.0

    def normalized_joint_position(self, joint_position: np.ndarray) -> np.ndarray:
        return (joint_position - self.joint_mid) / self.joint_half_range

    def build_observation(
        self,
        *,
        joint_position: np.ndarray,
        joint_velocity: np.ndarray,
        roll: float,
        pitch: float,
    ) -> np.ndarray:
        joint_position = np.asarray(joint_position, dtype=np.float32)
        joint_velocity = np.asarray(joint_velocity, dtype=np.float32)
        return np.concatenate(
            [
                self.normalized_joint_position(joint_position),
                joint_velocity,
                self.previous_action,
                np.array(
                    [
                        self.estimated_base_height,
                        roll,
                        pitch,
                        self.estimated_forward_velocity,
                        self.target_forward_velocity,
                    ],
                    dtype=np.float32,
                ),
            ],
            dtype=np.float32,
        )

    def action_to_joint_targets(self, action: np.ndarray) -> np.ndarray:
        action = np.asarray(action, dtype=np.float32)
        clipped_action = np.clip(action, -1.0, 1.0)
        reference_targets = self.nominal_joint_targets + self.reference_gait_offsets()
        desired_residual = clipped_action * self.joint_half_range * self.action_scale

        if self.action_smoothing > 0.0:
            smoothed_residual = (
                self.action_smoothing * self.current_action_residual
                + (1.0 - self.action_smoothing) * desired_residual
            )
        else:
            smoothed_residual = desired_residual

        self.current_action_residual = smoothed_residual.astype(np.float32)
        joint_targets = reference_targets + self.current_action_residual
        self.previous_action = clipped_action.copy()
        self.advance_gait_phase()
        return np.clip(joint_targets, self.joint_lower, self.joint_upper).astype(np.float32)

    def reference_gait_offsets(self) -> np.ndarray:
        offsets = np.zeros(len(self.joint_names), dtype=np.float32)
        if not self.reference_gait_enabled:
            return offsets

        hip_amplitude = float(self.reference_gait_config.get('hip_pitch_amplitude', 0.0))
        hip_bias = float(self.reference_gait_config.get('hip_pitch_bias', 0.0))
        knee_amplitude = float(self.reference_gait_config.get('knee_pitch_amplitude', 0.0))
        knee_bias = float(self.reference_gait_config.get('knee_pitch_bias', 0.0))
        swing_knee_scale = float(self.reference_gait_config.get('swing_knee_scale', 1.0))
        ankle_amplitude = float(self.reference_gait_config.get('ankle_pitch_amplitude', 0.0))
        ankle_bias = float(self.reference_gait_config.get('ankle_pitch_bias', 0.0))
        push_off_ankle_scale = float(self.reference_gait_config.get('push_off_ankle_scale', 0.0))
        stance_ratio = float(np.clip(self.reference_gait_config.get('stance_ratio', 0.62), 0.05, 0.95))

        def gait_profile(phase: float) -> tuple[float, float, float]:
            phase = phase % 1.0
            if phase < stance_ratio:
                stance_progress = phase / stance_ratio
                hip_profile = 1.0 - 2.0 * stance_progress
                knee_profile = 0.18 * math.sin(math.pi * stance_progress)
                push_off_progress = max((stance_progress - 0.60) / 0.40, 0.0)
                ankle_profile = -0.70 * hip_profile + push_off_ankle_scale * push_off_progress
            else:
                swing_progress = (phase - stance_ratio) / (1.0 - stance_ratio)
                hip_profile = -1.0 + 2.0 * swing_progress
                knee_profile = swing_knee_scale * math.sin(math.pi * swing_progress)
                ankle_profile = -0.45 * hip_profile - 0.25 * math.sin(math.pi * swing_progress)
            return hip_profile, knee_profile, ankle_profile

        left_hip, left_knee, left_ankle = gait_profile(self.gait_phase)
        right_hip, right_knee, right_ankle = gait_profile(self.gait_phase + 0.5)

        for index, joint_name in enumerate(self.joint_names):
            if joint_name == 'left_hip_pitch_joint':
                offsets[index] = hip_bias + hip_amplitude * left_hip
            elif joint_name == 'right_hip_pitch_joint':
                offsets[index] = hip_bias + hip_amplitude * right_hip
            elif joint_name == 'left_knee_pitch_joint':
                offsets[index] = knee_bias + knee_amplitude * left_knee
            elif joint_name == 'right_knee_pitch_joint':
                offsets[index] = knee_bias + knee_amplitude * right_knee
            elif joint_name == 'left_ankle_pitch_joint':
                offsets[index] = ankle_bias + ankle_amplitude * left_ankle
            elif joint_name == 'right_ankle_pitch_joint':
                offsets[index] = ankle_bias + ankle_amplitude * right_ankle

        return offsets

    def advance_gait_phase(self) -> None:
        if self.reference_gait_enabled:
            self.gait_phase = (self.gait_phase + self.control_dt / self.reference_gait_period) % 1.0
