# Copyright (c) 2025, Yutang Li, SIAT (yt.li2@siat.ac.cn)
# SPDX-License-Identifier: BSD-3-Clause

"""Chassis (mobile base) control via XR joystick."""

import torch


class ChassisController:
    """Converts left joystick input into wheel velocities and direct base velocity commands.

    Args:
        base_speed: Max wheel speed (rad/s) at full joystick.
        base_turn: Max wheel differential (rad/s) at full joystick.
        drive_speed: Max robot linear speed (m/s) for direct base control.
        drive_turn: Max robot yaw rate (rad/s) for direct base control.
    """

    def __init__(self, base_speed: float = 5.0, base_turn: float = 2.0,
                 drive_speed: float = 0.5, drive_turn: float = 1.5):
        self.base_speed = base_speed
        self.base_turn = base_turn
        self.drive_speed = drive_speed
        self.drive_turn = drive_turn

    def get_commands(self, xr_client) -> tuple[torch.Tensor, float, float]:
        """Read left joystick and return (wheel_cmd_4D, v_fwd_m_s, omega_rad_s)."""
        try:
            joy = xr_client.get_joystick("left")
            jy = float(joy[1])
            jx = float(joy[0])
        except Exception:
            return torch.zeros(4), 0.0, 0.0

        v_w = jy * self.base_speed
        omega_w = jx * self.base_turn
        right_vel = v_w - omega_w
        left_vel = v_w + omega_w
        wheel_cmd = torch.tensor([right_vel, left_vel, left_vel, right_vel], dtype=torch.float32)

        v_fwd = jy * self.drive_speed
        omega = -jx * self.drive_turn  # left push → positive yaw = left turn
        return wheel_cmd, v_fwd, omega

    @staticmethod
    def apply_base_velocity(robot, v_fwd: float, omega: float, num_envs: int, device) -> None:
        """Directly set robot root velocity to bypass isotropic friction for skid-steer turning."""
        if abs(v_fwd) < 1e-4 and abs(omega) < 1e-4:
            return
        rq = robot.data.root_quat_w  # [N, 4] wxyz
        w_q, x_q, y_q, z_q = rq[:, 0], rq[:, 1], rq[:, 2], rq[:, 3]
        fwd_x = 1.0 - 2.0 * (y_q * y_q + z_q * z_q)
        fwd_y = 2.0 * (x_q * y_q + w_q * z_q)
        vel = torch.zeros(num_envs, 6, device=device)
        vel[:, 0] = v_fwd * fwd_x
        vel[:, 1] = v_fwd * fwd_y
        vel[:, 5] = omega
        robot.write_root_velocity_to_sim(vel)