from __future__ import annotations

from collections.abc import Sequence
from dataclasses import MISSING

import torch

import isaaclab.utils.math as math_utils
from isaaclab.envs.mdp.commands import UniformVelocityCommand, UniformVelocityCommandCfg
from isaaclab.utils import configclass


class BodyVelocityCommand(UniformVelocityCommand):
    """Velocity command whose metrics and markers use a configured robot body."""

    cfg: BodyVelocityCommandCfg

    def __init__(self, cfg: BodyVelocityCommandCfg, env):
        if not 0.0 <= cfg.rel_straight_envs <= 1.0:
            raise ValueError(f"rel_straight_envs must be in [0, 1], got {cfg.rel_straight_envs}.")
        super().__init__(cfg, env)
        body_ids, body_names = self.robot.find_bodies(cfg.body_name, preserve_order=True)
        if len(body_ids) != 1:
            raise ValueError(
                f"BodyVelocityCommand requires exactly one body matching {cfg.body_name!r}; "
                f"found {body_names}."
            )
        self.body_id = body_ids[0]

    def _resample_command(self, env_ids: Sequence[int]):
        """Sample commands, with a configurable fraction reserved for straight walking."""
        env_ids = torch.as_tensor(env_ids, device=self.device, dtype=torch.long)
        super()._resample_command(env_ids)

        if self.cfg.rel_straight_envs <= 0.0 or env_ids.numel() == 0:
            return

        straight_mask = torch.rand(env_ids.numel(), device=self.device) <= self.cfg.rel_straight_envs
        straight_env_ids = env_ids[straight_mask]
        self.vel_command_b[straight_env_ids, 1:] = 0.0
        self.is_heading_env[straight_env_ids] = False

    def _body_heading_quat_w(self) -> torch.Tensor:
        body_quat_w = self.robot.data.body_quat_w[:, self.body_id, :]
        _, _, yaw = math_utils.euler_xyz_from_quat(body_quat_w)
        zeros = torch.zeros_like(yaw)
        return math_utils.quat_from_euler_xyz(zeros, zeros, yaw - self.cfg.heading_yaw_offset)

    def _body_lin_vel_heading(self) -> torch.Tensor:
        body_lin_vel_w = self.robot.data.body_lin_vel_w[:, self.body_id, :]
        return math_utils.quat_apply_inverse(self._body_heading_quat_w(), body_lin_vel_w)

    def _update_metrics(self):
        max_command_time = self.cfg.resampling_time_range[1]
        max_command_step = max_command_time / self._env.step_dt
        body_lin_vel_heading = self._body_lin_vel_heading()
        body_yaw_rate = self.robot.data.body_ang_vel_w[:, self.body_id, 2]
        self.metrics["error_vel_xy"] += (
            torch.norm(self.vel_command_b[:, :2] - body_lin_vel_heading[:, :2], dim=-1) / max_command_step
        )
        self.metrics["error_vel_yaw"] += (
            torch.abs(self.vel_command_b[:, 2] - body_yaw_rate) / max_command_step
        )

    def _update_command(self):
        if self.cfg.heading_command:
            env_ids = self.is_heading_env.nonzero(as_tuple=False).flatten()
            if len(env_ids) > 0:
                body_quat_w = self.robot.data.body_quat_w[env_ids, self.body_id, :]
                _, _, body_yaw = math_utils.euler_xyz_from_quat(body_quat_w)
                body_heading = body_yaw - self.cfg.heading_yaw_offset
                heading_error = math_utils.wrap_to_pi(self.heading_target[env_ids] - body_heading)
                self.vel_command_b[env_ids, 2] = torch.clip(
                    self.cfg.heading_control_stiffness * heading_error,
                    min=self.cfg.ranges.ang_vel_z[0],
                    max=self.cfg.ranges.ang_vel_z[1],
                )

        standing_env_ids = self.is_standing_env.nonzero(as_tuple=False).flatten()
        self.vel_command_b[standing_env_ids, :] = 0.0

    def _debug_vis_callback(self, _event):
        if not self.robot.is_initialized:
            return

        marker_pos_w = self.robot.data.body_pos_w[:, self.body_id, :].clone()
        marker_pos_w[:, 2] += self.cfg.marker_height_offset
        desired_pos_w = marker_pos_w.clone()
        actual_pos_w = marker_pos_w.clone()
        desired_pos_w[:, 2] += 0.5 * self.cfg.marker_vertical_separation
        actual_pos_w[:, 2] -= 0.5 * self.cfg.marker_vertical_separation
        desired_scale, desired_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2])
        actual_scale, actual_quat = self._resolve_xy_velocity_to_arrow(self._body_lin_vel_heading()[:, :2])
        self.goal_vel_visualizer.visualize(desired_pos_w, desired_quat, desired_scale)
        self.current_vel_visualizer.visualize(actual_pos_w, actual_quat, actual_scale)

    def _resolve_xy_velocity_to_arrow(self, xy_velocity: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        default_scale = self.goal_vel_visualizer.cfg.markers["arrow"].scale
        arrow_scale = torch.tensor(default_scale, device=self.device).repeat(xy_velocity.shape[0], 1)
        arrow_scale[:, 0] *= torch.linalg.norm(xy_velocity, dim=1) * 3.0

        heading_angle = torch.atan2(xy_velocity[:, 1], xy_velocity[:, 0])
        zeros = torch.zeros_like(heading_angle)
        arrow_quat = math_utils.quat_from_euler_xyz(zeros, zeros, heading_angle)
        arrow_quat = math_utils.quat_mul(self._body_heading_quat_w(), arrow_quat)
        return arrow_scale, arrow_quat


@configclass
class BodyVelocityCommandCfg(UniformVelocityCommandCfg):
    """Configuration for body-referenced planar velocity commands."""

    class_type: type = BodyVelocityCommand

    body_name: str = MISSING
    """Body used for velocity metrics and command visualization."""

    heading_yaw_offset: float = 0.0
    """Yaw offset from the body's URDF frame to the locomotion heading frame."""

    rel_straight_envs: float = 0.0
    """Fraction of sampled environments with zero lateral and yaw commands."""

    marker_height_offset: float = 0.35
    """Vertical marker offset from the configured body origin in meters."""

    marker_vertical_separation: float = 0.08
    """Vertical separation between desired and measured velocity arrows in meters."""