"""Unitree Go2-W 原地站立平衡控制器（纯 MuJoCo/Python 版本）。

兼容两种模型：

* unitree_mujoco/unitree_robots/go2w/go2w.xml；
* Web 平台由 unitree_ros/go2w_description.urdf 转换出的浮动基座 MJCF。

控制策略参考 unitree_mujoco/example/python/stand_go2.py：站立关节角、kp=50、
kd=3.5，并在浏览器中直接计算 motor torque。姿态由基座6轴 IMU 的角速度与
重力加速度互补滤波获得，不包含 DDS、CRC、实机电机模式或 sim2real 逻辑。
"""

import math

NAME = "Unitree Go2-W 原地平衡控制"
CONTROL_HZ = 200

LEG_NAMES = ("FL", "FR", "RL", "RR")

# 来自 unitree_mujoco/example/python/stand_go2.py 的 stand_up_joint_pos。
HIP_TARGET = {"FL": 0.00571868, "FR": -0.00571868,
              "RL": 0.00571868, "RR": -0.00571868}
NOMINAL = {"thigh": 0.608813, "calf": -1.21763}
HIP_DOWN = {"FL": 0.0473455, "FR": -0.0473455,
            "RL": 0.0473455, "RR": -0.0473455}
STAND_DOWN = {"thigh": 1.22187, "calf": -2.44375}

# SDK 示例使用 kp=50、kd=3.5。力矩上限按 Go2-W URDF/MJCF 保守取值。
KP_FINAL = 50.0
KP_INITIAL = 20.0
KD = 3.5
EFFORT_LIMIT = {"hip": 23.7, "thigh": 23.7, "calf": 35.55}
RAMP_SECONDS = 1.2
WHEEL_DAMPING = 0.12
MAX_WHEEL_TORQUE = 6.0


def _clamp(value, lower, upper):
    return max(lower, min(upper, value))


def _accel_attitude(acceleration):
    """由 Body 局部系中的重力方向估计横滚角和俯仰角。"""
    ax, ay, az = acceleration
    return math.atan2(ay, az), math.atan2(-ax, math.sqrt(ay * ay + az * az))


def _resolve(api, kind, names):
    """依次尝试官方 MJCF 与 URDF 转换模型的命名。"""
    resolver = getattr(api, kind)
    last_error = None
    for name in names:
        try:
            return resolver(name)
        except Exception as error:
            last_error = error
    raise RuntimeError(f"无法解析 {kind}，候选名称：{', '.join(names)}") from last_error


def init(api):
    legs = []
    for prefix in LEG_NAMES:
        joints = {}
        actuators = {}
        for part in ("hip", "thigh", "calf"):
            joint_name = f"{prefix}_{part}_joint"
            joints[part] = _resolve(api, "joint", (joint_name,))
            actuators[part] = _resolve(
                api, "actuator", (f"{prefix}_{part}", f"{joint_name}_motor")
            )

        wheel_joint = _resolve(
            api, "joint", (f"{prefix}_wheel_joint", f"{prefix}_foot_joint")
        )
        wheel_actuator = _resolve(
            api,
            "actuator",
            (f"{prefix}_wheel", f"{prefix}_wheel_joint_motor", f"{prefix}_foot_joint_motor"),
        )
        legs.append({
            "name": prefix,
            "joints": joints,
            "actuators": actuators,
            "wheel_joint": wheel_joint,
            "wheel_actuator": wheel_actuator,
            "side": 1.0 if prefix.endswith("L") else -1.0,
            "fore": 1.0 if prefix.startswith("F") else -1.0,
        })

    return {
        "base": _resolve(api, "body", ("base_link", "base")),
        "imu_gyro": _resolve(api, "sensor", ("imu_gyro", "__platform_imu_gyro__")),
        "imu_acc": _resolve(api, "sensor", ("imu_acc", "__platform_imu_acc__")),
        "legs": legs,
        "started": False,
        "start_time": 0.0,
        "estimated_roll": 0.0,
        "estimated_pitch": 0.0,
        "filtered_droll": 0.0,
        "filtered_dpitch": 0.0,
        "unstable_duration": 0.0,
    }


def _initialize(ctx, state, acceleration):
    state["started"] = True
    state["start_time"] = ctx.time
    state["estimated_roll"], state["estimated_pitch"] = _accel_attitude(acceleration)


def step(ctx, state):
    gyro = ctx.sensor(state["imu_gyro"])
    acceleration = ctx.sensor(state["imu_acc"])
    base_position = ctx.body_position(state["base"])
    if not state["started"]:
        _initialize(ctx, state, acceleration)

    dt = max(1.0e-4, ctx.dt)
    rate_alpha = 0.25
    state["filtered_droll"] += rate_alpha * (gyro[0] - state["filtered_droll"])
    state["filtered_dpitch"] += rate_alpha * (gyro[1] - state["filtered_dpitch"])
    accel_roll, accel_pitch = _accel_attitude(acceleration)
    fusion = 0.015
    state["estimated_roll"] = (1.0 - fusion) * (
        state["estimated_roll"] + state["filtered_droll"] * dt
    ) + fusion * accel_roll
    state["estimated_pitch"] = (1.0 - fusion) * (
        state["estimated_pitch"] + state["filtered_dpitch"] * dt
    ) + fusion * accel_pitch
    roll, pitch = state["estimated_roll"], state["estimated_pitch"]

    elapsed = max(0.0, ctx.time - state["start_time"])
    ramp = 1.0 if elapsed >= 3.0 else math.tanh(elapsed / RAMP_SECONDS)
    kp = KP_INITIAL + (KP_FINAL - KP_INITIAL) * ramp

    # SDK 示例的核心：12 个腿关节平滑进入站立姿态并保持 PD 闭环。
    for leg in state["legs"]:
        for part in ("hip", "thigh", "calf"):
            up = HIP_TARGET[leg["name"]] if part == "hip" else NOMINAL[part]
            down = HIP_DOWN[leg["name"]] if part == "hip" else STAND_DOWN[part]
            desired = down + ramp * (up - down)
            if part == "calf" and ramp > 0.8:
                # 通过左右/前后轮腿长度差调平车身：高侧缩短，低侧伸长。
                roll_term = 0.34 * roll + 0.025 * state["filtered_droll"]
                pitch_term = 0.30 * pitch + 0.022 * state["filtered_dpitch"]
                desired -= leg["side"] * roll_term
                desired += leg["fore"] * pitch_term
            target = desired
            position = ctx.qpos(leg["joints"][part])
            velocity = ctx.qvel(leg["joints"][part])
            torque = kp * (target - position) - KD * velocity
            limit = EFFORT_LIMIT[part]
            ctx.set_control(leg["actuators"][part], _clamp(torque, -limit, limit))

    # Go2-W 四轮保持自由滚动，只施加温和阻尼以抑制无指令漂移。
    wheel_velocities = [ctx.qvel(leg["wheel_joint"]) for leg in state["legs"]]
    average_velocity = sum(wheel_velocities) / len(wheel_velocities)
    wheel_torque = 1.4 * pitch + 0.10 * state["filtered_dpitch"]
    wheel_torque -= WHEEL_DAMPING * average_velocity
    wheel_torque = _clamp(wheel_torque, -MAX_WHEEL_TORQUE, MAX_WHEEL_TORQUE)
    for leg in state["legs"]:
        ctx.set_control(leg["wheel_actuator"], wheel_torque)

    # 持续失稳时主动停止，避免倒地后控制器继续输出饱和力矩。
    unstable = ramp > 0.95 and (
        base_position[2] < 0.16 or abs(roll) > 1.0 or abs(pitch) > 1.0
    )
    state["unstable_duration"] = state["unstable_duration"] + dt if unstable else 0.0
    if state["unstable_duration"] > 0.35:
        raise RuntimeError(
            f"Go2-W 已失稳：z={base_position[2]:.3f} m, "
            f"roll={roll:.3f} rad, pitch={pitch:.3f} rad；请重置后检查模型接触参数"
        )


def reset(state):
    state["started"] = False
    state["estimated_roll"] = 0.0
    state["estimated_pitch"] = 0.0
    state["filtered_droll"] = 0.0
    state["filtered_dpitch"] = 0.0
    state["unstable_duration"] = 0.0


def dispose(state):
    pass