#!/usr/bin/env python3
"""ROS 2 MuJoCo bridge for the original M20 and the 8-DoF FireAnt model."""

from __future__ import annotations

import argparse
import time
from pathlib import Path

import mujoco
import mujoco.viewer
import numpy as np
import rclpy
from builtin_interfaces.msg import Time
from drdds.msg import (
    ImuData,
    ImuDataValue,
    JointData,
    JointsData,
    JointsDataCmd,
    JointsDataValue,
    MetaType,
)
from rclpy.node import Node


CURRENT_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = CURRENT_DIR.parents[2]
MESSAGE_JOINT_COUNT = 16
DT = 0.001
STATE_PUBLISH_DECIMATION = 5  # 200 Hz, matching the C++ controller loop.

FIRE_ANT_JOINTS = (
    "fl_knee_joint", "fl_wheel_joint",
    "fr_knee_joint", "fr_wheel_joint",
    "hl_knee_joint", "hl_wheel_joint",
    "hr_knee_joint", "hr_wheel_joint",
)
FIRE_ANT_HOME = np.array([0.5, 0.0, 0.5, 0.0, -0.5, 0.0, -0.5, 0.0], dtype=np.float64)

M20_DIR = np.array(
    [1, 1, -1, 1, 1, -1, 1, -1, -1, 1, -1, 1, -1, -1, 1, -1],
    dtype=np.float64,
)
M20_OFFSET = np.deg2rad(
    np.array([-25, 229, 160, 0, 25, -131, -200, 0, -25, -229, -160, 0, 25, 131, 200, 0])
)
M20_HOME = np.array(
    [-0.438, -1.16, 2.76, 0, 0.438, -1.16, 2.76, 0,
     -0.438, 1.16, -2.76, 0, 0.438, 1.16, -2.76, 0],
    dtype=np.float64,
)


def model_path(robot: str, terrain: str) -> Path:
    if robot == "fire_ant":
        name = "fire_ant_rough.xml" if terrain == "rough" else "fire_ant.xml"
        return PROJECT_ROOT / "M20_description" / "fire_ant_mjcf" / "mjcf" / name
    name = "M20_stair.xml" if terrain == "rough" else "M20.xml"
    return PROJECT_ROOT / "M20_description" / "m20_mjcf" / "mjcf" / name


def quaternion_to_euler(q: np.ndarray) -> np.ndarray:
    """Convert a MuJoCo wxyz quaternion to roll, pitch and yaw."""
    w, x, y, z = q
    roll = np.arctan2(2.0 * (w * x + y * z), 1.0 - 2.0 * (x * x + y * y))
    pitch = np.arcsin(np.clip(2.0 * (w * y - z * x), -1.0, 1.0))
    yaw = np.arctan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z))
    return np.array([roll, pitch, yaw], dtype=np.float64)


class MuJoCoSimulationNode(Node):
    def __init__(
        self,
        robot: str,
        xml_path: Path,
        use_viewer: bool,
        randomize: bool,
        seed: int,
        realtime_factor: float,
    ) -> None:
        super().__init__("mujoco_simulation")
        if not xml_path.is_file():
            raise FileNotFoundError(f"Cannot find MJCF: {xml_path}")

        self.robot = robot
        self.model = mujoco.MjModel.from_xml_path(str(xml_path))
        self.model.opt.timestep = DT
        self.data = mujoco.MjData(self.model)
        self.realtime_factor = realtime_factor

        self.dof_num = self.model.nu
        expected_dof = 8 if robot == "fire_ant" else 16
        if self.dof_num != expected_dof:
            raise RuntimeError(f"{robot} expected {expected_dof} actuators, model has {self.dof_num}")

        self.joint_names = list(FIRE_ANT_JOINTS) if robot == "fire_ant" else [
            mujoco.mj_id2name(
                self.model,
                mujoco.mjtObj.mjOBJ_JOINT,
                int(self.model.actuator_trnid[i, 0]),
            )
            for i in range(self.dof_num)
        ]
        self.joint_ids = np.array(
            [mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, name) for name in self.joint_names],
            dtype=np.int32,
        )
        if np.any(self.joint_ids < 0):
            raise RuntimeError(f"Missing joints in MJCF: {self.joint_names}")
        self.qpos_adr = self.model.jnt_qposadr[self.joint_ids]
        self.dof_adr = self.model.jnt_dofadr[self.joint_ids]
        self.actuator_ids = np.array(
            [
                next(i for i in range(self.model.nu) if int(self.model.actuator_trnid[i, 0]) == joint_id)
                for joint_id in self.joint_ids
            ],
            dtype=np.int32,
        )

        self.joint_dir = np.ones(self.dof_num) if robot == "fire_ant" else M20_DIR.copy()
        self.pos_offset = np.zeros(self.dof_num) if robot == "fire_ant" else M20_OFFSET.copy()
        self.kp_cmd = np.zeros(self.dof_num)
        self.kd_cmd = np.zeros(self.dof_num)
        self.pos_cmd = np.zeros(self.dof_num)
        self.vel_cmd = np.zeros(self.dof_num)
        self.tau_ff = np.zeros(self.dof_num)
        self.applied_tau = np.zeros(self.dof_num)
        self.velocity_limit = (
            np.array([17.8024, 109.956] * 4)
            if robot == "fire_ant"
            else np.array([45.0, 22.4, 22.4, 100.0] * 4)
        )
        self.kp_scale = np.ones(self.dof_num)
        self.kd_scale = np.ones(self.dof_num)
        self.timestamp = 0.0

        self.base_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, "base_link")
        self.wheel_body_ids: list[int] = []
        self.knee_body_ids: list[int] = []
        if robot == "fire_ant":
            self.wheel_body_ids = [self._body_id(f"{leg}_wheel") for leg in ("fl", "fr", "hl", "hr")]
            self.knee_body_ids = [self._body_id(f"{leg}_knee") for leg in ("fl", "fr", "hl", "hr")]
            if randomize:
                self._apply_training_randomization(seed)

        self._set_initial_pose()
        self.imu_pub = self.create_publisher(ImuData, "/IMU_DATA", 200)
        self.joints_pub = self.create_publisher(JointsData, "/JOINTS_DATA", 200)
        self.cmd_sub = self.create_subscription(JointsDataCmd, "/JOINTS_CMD", self._cmd_callback, 50)
        self.viewer = mujoco.viewer.launch_passive(self.model, self.data) if use_viewer else None

        mode = "randomized" if randomize else "nominal"
        self.get_logger().info(
            f"Loaded {robot} ({mode}) from {xml_path}; dof={self.dof_num}, dt={DT:.4f}s"
        )

    def _body_id(self, name: str) -> int:
        body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, name)
        if body_id < 0:
            raise RuntimeError(f"Missing body in MJCF: {name}")
        return body_id

    def _apply_training_randomization(self, seed: int) -> None:
        """Apply one reproducible sample from the training-time domain ranges."""
        rng = np.random.default_rng(seed)
        original_mass = self.model.body_mass.copy()
        for body_id in [*self.knee_body_ids, *self.wheel_body_ids]:
            mass_scale = rng.uniform(0.85, 1.15)
            inertia_scale = mass_scale * rng.uniform(0.85, 1.15)
            self.model.body_mass[body_id] *= mass_scale
            self.model.body_inertia[body_id] *= inertia_scale

        base_mass = max(0.1, original_mass[self.base_id] + rng.uniform(-1.0, 3.0))
        base_mass_scale = base_mass / original_mass[self.base_id]
        self.model.body_mass[self.base_id] = base_mass
        self.model.body_inertia[self.base_id] *= base_mass_scale * rng.uniform(0.85, 1.15)
        self.model.body_ipos[self.base_id] += rng.uniform((-0.03, -0.03, -0.02), (0.03, 0.03, 0.02))

        for body_id in self.wheel_body_ids:
            geom_start = self.model.body_geomadr[body_id]
            geom_count = self.model.body_geomnum[body_id]
            self.model.geom_friction[geom_start:geom_start + geom_count, 0] = rng.uniform(0.3, 1.0)
        self.kp_scale = rng.uniform(0.85, 1.15, self.dof_num)
        self.kd_scale = rng.uniform(0.85, 1.15, self.dof_num)

    def _set_initial_pose(self) -> None:
        self.data.qpos[:] = 0.0
        self.data.qvel[:] = 0.0
        self.data.qpos[0:3] = (0.0, 0.0, 0.306 if self.robot == "fire_ant" else 0.2)
        self.data.qpos[3:7] = (1.0, 0.0, 0.0, 0.0)
        self.data.qpos[self.qpos_adr] = FIRE_ANT_HOME if self.robot == "fire_ant" else M20_HOME
        mujoco.mj_forward(self.model, self.data)

    def _cmd_callback(self, msg: JointsDataCmd) -> None:
        if len(msg.data.joints_data) < self.dof_num:
            self.get_logger().warn(
                f"Ignoring short JointsDataCmd ({len(msg.data.joints_data)} < {self.dof_num})"
            )
            return
        pub_pos = np.empty(self.dof_num)
        pub_vel = np.empty(self.dof_num)
        for i in range(self.dof_num):
            command = msg.data.joints_data[i]
            self.kp_cmd[i] = command.kp
            self.kd_cmd[i] = command.kd
            pub_pos[i] = command.position
            pub_vel[i] = command.velocity
            self.tau_ff[i] = command.torque
        self.pos_cmd = pub_pos * self.joint_dir + self.pos_offset
        self.vel_cmd = pub_vel * self.joint_dir

    def _joint_state(self) -> tuple[np.ndarray, np.ndarray]:
        return self.data.qpos[self.qpos_adr].copy(), self.data.qvel[self.dof_adr].copy()

    def _apply_joint_torque(self) -> None:
        q, dq = self._joint_state()
        torque = (
            self.kp_cmd * self.kp_scale * (self.pos_cmd - q)
            + self.kd_cmd * self.kd_scale * (self.vel_cmd - dq)
            + self.tau_ff
        )
        ctrl_range = self.model.actuator_ctrlrange[self.actuator_ids]
        torque = np.clip(torque, ctrl_range[:, 0], ctrl_range[:, 1])
        # The MuJoCo motor has a torque limit but no built-in speed limit. Prevent
        # positive work beyond the datasheet/training actuator velocity limit.
        overspeed_accelerating = (np.abs(dq) >= self.velocity_limit) & (torque * dq > 0.0)
        torque[overspeed_accelerating] = 0.0
        self.data.ctrl[self.actuator_ids] = torque

    def _body_linear_velocity(self) -> np.ndarray:
        # Free-joint translational qvel is in world axes; the policy uses body axes.
        rotation_world_from_body = self.data.xmat[self.base_id].reshape(3, 3)
        return rotation_world_from_body.T @ self.data.qvel[0:3]

    def _contact_force_norms(self, body_ids: list[int]) -> np.ndarray:
        # cfrc_ext stores [torque, force]; norm is frame invariant.
        return np.linalg.norm(self.data.cfrc_ext[body_ids, 3:6], axis=1)

    def _make_stamp(self) -> Time:
        stamp = Time()
        stamp.sec = int(self.timestamp)
        stamp.nanosec = int((self.timestamp - stamp.sec) * 1e9)
        return stamp

    def _publish_robot_state(self) -> None:
        q_world = self.data.sensordata[0:4]
        rpy_deg = np.rad2deg(quaternion_to_euler(q_world))
        body_acc = self.data.sensordata[4:7]
        angular_velocity_body = self.data.sensordata[7:10]

        imu_msg = ImuData()
        imu_msg.header = MetaType()
        imu_msg.header.frame_id = 0
        imu_msg.header.stamp = self._make_stamp()
        imu_msg.data = ImuDataValue()
        imu_msg.data.roll, imu_msg.data.pitch, imu_msg.data.yaw = map(float, rpy_deg)
        imu_msg.data.omega_x, imu_msg.data.omega_y, imu_msg.data.omega_z = map(
            float, angular_velocity_body
        )
        imu_msg.data.acc_x, imu_msg.data.acc_y, imu_msg.data.acc_z = map(float, body_acc)
        self.imu_pub.publish(imu_msg)

        q, dq = self._joint_state()
        pub_pos = (q - self.pos_offset) * self.joint_dir
        pub_vel = dq * self.joint_dir
        self.applied_tau = self.data.actuator_force[self.actuator_ids] * self.joint_dir

        joints_msg = JointsData()
        joints_msg.header = MetaType()
        joints_msg.header.frame_id = 0
        joints_msg.header.stamp = self._make_stamp()
        joints_msg.data = JointsDataValue()
        joints_msg.data.joints_data = [JointData() for _ in range(MESSAGE_JOINT_COUNT)]
        for i in range(self.dof_num):
            joint = joints_msg.data.joints_data[i]
            joint.name = [32, 32, 32, 32]
            joint.data_id = 0
            joint.status_word = 1
            joint.position = float(pub_pos[i])
            joint.velocity = float(pub_vel[i])
            joint.torque = float(self.applied_tau[i])
            joint.motion_temp = 40.0
            joint.driver_temp = 45.0

        if self.robot == "fire_ant":
            # Sim-only extension inside the eight unused fields of the fixed 16-joint DDS message.
            wheel_force = self._contact_force_norms(self.wheel_body_ids)
            knee_force = self._contact_force_norms(self.knee_body_ids)
            base_force = self._contact_force_norms([self.base_id])[0]
            base_velocity = self._body_linear_velocity()
            for i in range(4):
                joints_msg.data.joints_data[8 + i].position = float(wheel_force[i])
                joints_msg.data.joints_data[12 + i].position = float(knee_force[i])
            for i in range(3):
                joints_msg.data.joints_data[8 + i].torque = float(base_velocity[i])
            joints_msg.data.joints_data[12].velocity = float(base_force)

        self.joints_pub.publish(joints_msg)

    def start(self) -> None:
        step = 0
        next_step_time = time.perf_counter()
        while rclpy.ok() and (self.viewer is None or self.viewer.is_running()):
            rclpy.spin_once(self, timeout_sec=0.0)
            self._apply_joint_torque()
            mujoco.mj_step(self.model, self.data)
            mujoco.mj_rnePostConstraint(self.model, self.data)
            step += 1
            self.timestamp = step * DT
            if step % STATE_PUBLISH_DECIMATION == 0:
                self._publish_robot_state()
            if self.viewer is not None and step % 20 == 0:
                self.viewer.sync()

            if self.realtime_factor > 0.0:
                next_step_time += DT / self.realtime_factor
                remaining = next_step_time - time.perf_counter()
                if remaining > 0:
                    time.sleep(remaining)
                elif remaining < -0.1:
                    next_step_time = time.perf_counter()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--robot", choices=("fire_ant", "m20"), default="fire_ant")
    parser.add_argument("--terrain", choices=("flat", "rough"), default="rough")
    parser.add_argument("--xml", type=Path, help="Override the selected MJCF path")
    parser.add_argument("--headless", action="store_true", help="Run without the MuJoCo viewer")
    parser.add_argument("--randomize", action="store_true", help="Sample training mass/friction/gain ranges once")
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--realtime-factor", type=float, default=1.0, help="0 runs as fast as possible")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    xml_path = args.xml.resolve() if args.xml else model_path(args.robot, args.terrain).resolve()
    rclpy.init()
    node = MuJoCoSimulationNode(
        robot=args.robot,
        xml_path=xml_path,
        use_viewer=not args.headless,
        randomize=args.randomize,
        seed=args.seed,
        realtime_factor=args.realtime_factor,
    )
    try:
        node.start()
    except KeyboardInterrupt:
        pass
    finally:
        node.destroy_node()
        if rclpy.ok():
            rclpy.shutdown()


if __name__ == "__main__":
    np.set_printoptions(precision=4, suppress=True)
    main()