# Copyright (c) 2022-2026, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

# Flag for pyright to ignore type errors in this file.
# pyright: reportPrivateUsage=false

from __future__ import annotations

import logging
from collections.abc import Sequence
from typing import TYPE_CHECKING

import torch

import isaaclab.utils.string as string_utils
from isaaclab.assets.articulation import Articulation

from isaaclab_contrib.actuators import Thruster
from isaaclab_contrib.utils.types import MultiRotorActions

from .multirotor_data import MultirotorData

if TYPE_CHECKING:
    from .multirotor_cfg import MultirotorCfg

# import logger
logger = logging.getLogger(__name__)


class Multirotor(Articulation):
    """A multirotor articulation asset class.

    This class extends the base :class:`~isaaclab.assets.Articulation` class to support multirotor vehicles
    (such as quadcopters, hexacopters, and octocopters) with thruster actuators that apply forces
    at specific body locations. It is based on the implementation from :cite:t:`kulkarni2025aerialgym`.

    Unlike standard articulations that use joint-based control, multirotors are controlled through
    thrust forces generated by individual rotors/propellers. This class provides specialized functionality
    for managing multiple thruster actuators, computing combined wrenches from individual thrusts,
    and applying them to the multirotor's base link.

    Key Features:
        - **Thruster-based control**: Uses :class:`~isaaclab_contrib.actuators.Thruster` actuators
          instead of joint actuators for realistic rotor dynamics simulation.
        - **Force allocation**: Supports allocation matrices to convert individual thruster forces
          into combined body wrenches (forces and torques).
        - **Asymmetric dynamics**: Thruster actuators can model asymmetric rise/fall dynamics
          that reflect real motor behavior.
        - **Flexible configuration**: Supports arbitrary numbers and arrangements of thrusters
          through regex-based thruster naming patterns.

    Usage Example:
        .. code-block:: python

            import isaaclab.sim as sim_utils
            from isaaclab_contrib.assets import MultirotorCfg
            from isaaclab_contrib.actuators import ThrusterCfg

            # Define thruster actuator configuration
            thruster_cfg = ThrusterCfg(
                thruster_names_expr=["rotor_[0-3]"],  # Match rotors 0-3
                thrust_range=(0.0, 10.0),  # Min and max thrust in Newtons
                rise_time_constant=0.1,  # Time constant for thrust increase
                fall_time_constant=0.2,  # Time constant for thrust decrease
            )

            # Create multirotor configuration
            multirotor_cfg = MultirotorCfg(
                prim_path="/World/envs/env_.*/Robot",
                spawn=sim_utils.UsdFileCfg(usd_path="path/to/quadcopter.usd"),
                actuators={"thrusters": thruster_cfg},
                allocation_matrix=[  # 6x4 matrix for quadcopter (6 DOF, 4 thrusters)
                    [1.0, 1.0, 1.0, 1.0],  # Total vertical force
                    [0.0, 0.0, 0.0, 0.0],  # Lateral force (x)
                    [0.0, 0.0, 0.0, 0.0],  # Lateral force (y)
                    [0.0, 0.1, 0.0, -0.1],  # Roll torque
                    [-0.1, 0.0, 0.1, 0.0],  # Pitch torque
                    [0.01, -0.01, 0.01, -0.01],  # Yaw torque
                ],
            )

            # Create the multirotor instance
            multirotor = multirotor_cfg.class_type(multirotor_cfg)

    .. note::
        The allocation matrix maps individual thruster forces to a 6D wrench (3 forces + 3 torques)
        applied to the base link. The matrix dimensions should be (6, num_thrusters).

    .. seealso::
        - :class:`~isaaclab.assets.Articulation`: Base articulation class
        - :class:`MultirotorCfg`: Configuration class for multirotors
        - :class:`MultirotorData`: Data container for multirotor state
        - :class:`~isaaclab_contrib.actuators.Thruster`: Thruster actuator model
    """

    cfg: MultirotorCfg
    """Configuration instance for the multirotor."""

    actuators: dict[str, Thruster]
    """Dictionary of thruster actuator instances for the multirotor.

    The keys are the actuator names and the values are the actuator instances. The actuator instances
    are initialized based on the actuator configurations specified in the :attr:`MultirotorCfg.actuators`
    attribute. They are used to compute the thruster commands during the :meth:`write_data_to_sim` function.
    """

    def __init__(self, cfg: MultirotorCfg):
        """Initialize the multirotor articulation.

        Args:
            cfg: A configuration instance.
        """
        super().__init__(cfg)

    """
    Properties
    """

    @property
    def thruster_names(self) -> list[str]:
        """Ordered names of thrusters in the multirotor.

        This property aggregates thruster names from all thruster actuator groups configured
        for the multirotor. The names are ordered according to their array indices, which is
        important for setting thrust targets and interpreting thruster data.

        Returns:
            A list of thruster names in order. Returns an empty list if actuators are not yet initialized.

        Raises:
            ValueError: If a non-thruster actuator is found in the multirotor actuators.
        """
        if not hasattr(self, "actuators") or not self.actuators:
            return []

        thruster_names = []
        for actuator in self.actuators.values():
            if hasattr(actuator, "thruster_names"):
                thruster_names.extend(actuator.thruster_names)
            else:
                raise ValueError("Non thruster actuator found in multirotor actuators. Not supported at the moment.")

        return thruster_names

    @property
    def num_thrusters(self) -> int:
        """Number of thrusters in the multirotor.

        Returns:
            Total number of thrusters across all actuator groups.
        """
        return len(self.thruster_names)

    @property
    def allocation_matrix(self) -> torch.Tensor:
        """Allocation matrix for control allocation.

        The allocation matrix maps individual thruster forces to a 6D wrench vector
        (3 forces + 3 torques) applied to the base link. This allows converting
        per-thruster commands into the resulting body-frame forces and moments.

        The matrix has shape (6, num_thrusters), where:
            - Rows 0-2: Force contributions in body frame (Fx, Fy, Fz)
            - Rows 3-5: Torque contributions in body frame (Tx, Ty, Tz)

        Returns:
            Allocation matrix as a torch tensor on the device.
        """
        return torch.tensor(self.cfg.allocation_matrix, device=self.device, dtype=torch.float32)

    """
    Operations
    """

    def set_thrust_target(
        self,
        target: torch.Tensor,
        thruster_ids: Sequence[int] | slice | None = None,
        env_ids: Sequence[int] | None = None,
    ):
        """Set target thrust values for thrusters.

        This method sets the desired thrust values for specific thrusters in specific environments.
        The thrust targets are stored and later processed by the thruster actuator models during
        the :meth:`write_data_to_sim` call. The actuator models may apply dynamics (rise/fall times)
        and constraints (thrust limits) to these targets.

        Args:
            target: Target thrust values. Shape is (num_envs, num_thrusters) or (num_envs,).
                The values are typically in the same units as configured in the thruster actuator
                (e.g., Newtons for force, or revolutions per second for RPS).
            thruster_ids: Indices of thrusters to set. Defaults to None (all thrusters).
                Can be a sequence of integers, a slice, or None.
            env_ids: Environment indices to set. Defaults to None (all environments).
                Can be a sequence of integers or None.

        Example:
            .. code-block:: python

                # Set thrust for all thrusters in all environments
                multirotor.set_thrust_target(torch.ones(num_envs, 4) * 5.0)

                # Set thrust for specific thrusters
                multirotor.set_thrust_target(
                    torch.tensor([[5.0, 6.0]]),  # Different thrust for 2 thrusters
                    thruster_ids=[0, 2],  # Apply to thrusters 0 and 2
                    env_ids=[0],  # Only in environment 0
                )
        """
        # resolve indices
        if env_ids is None:
            env_ids = slice(None)
        if thruster_ids is None:
            thruster_ids = slice(None)

        # broadcast env_ids if needed to allow double indexing
        if env_ids != slice(None) and thruster_ids != slice(None):
            env_ids = env_ids[:, None]

        # set targets
        self._data.thrust_target[env_ids, thruster_ids] = target

    def reset(self, env_ids: Sequence[int] | None = None):
        """Reset the multirotor to default state.

        This method resets both the base articulation state (pose, velocities) and
        multirotor-specific state (thruster targets) to their default values as specified
        in the configuration.

        Args:
            env_ids: Environment indices to reset. Defaults to None (all environments).
                Can be a sequence of integers or None.

        Note:
            The default thruster state is set via the :attr:`MultirotorCfg.init_state.rps`
            configuration parameter.
        """
        # call parent reset
        super().reset(env_ids)

        # reset multirotor-specific data
        if env_ids is None:
            env_ids = self._ALL_INDICES
        elif not isinstance(env_ids, torch.Tensor):
            env_ids = torch.tensor(env_ids, dtype=torch.long, device=self.device)

        # reset thruster targets to default values
        if self._data.thrust_target is not None and self._data.default_thruster_rps is not None:
            self._data.thrust_target[env_ids] = self._data.default_thruster_rps[env_ids]

    def write_data_to_sim(self):
        """Write thrust and torque commands to the simulation.

        This method performs the following operations in sequence:

        1. **Apply actuator models**: Process thrust targets through thruster actuator models
           to compute actual thrust values considering dynamics (rise/fall times) and
           constraints (thrust limits).

        2. **Combine thrusts into wrench**: Use the allocation matrix to convert individual
           thruster forces into a combined 6D wrench (force + torque) vector.

        3. **Apply to simulation**: Apply the combined wrench to the base link of the multirotor
           in the PhysX simulation.

        This method should be called after setting thrust targets with :meth:`set_thrust_target`
        and before stepping the simulation.

        Note:
            This method overrides the base class implementation because multirotors use thrust-based
            control rather than joint-based control.
        """
        self._apply_actuator_model()
        # apply thruster forces at individual locations
        self._apply_combined_wrench()

    """
    Internal methods
    """

    def _initialize_impl(self):
        """Initialize the multirotor implementation."""
        # call parent initialization
        super()._initialize_impl()

        # Replace data container with MultirotorData
        self._data = MultirotorData(self.root_physx_view, self.device)

        # Create thruster buffers with correct size (SINGLE PHASE)
        self._create_thruster_buffers()
        # Process thruster configuration
        self._process_thruster_cfg()
        # Process configuration
        self._process_cfg()
        # Update the robot data
        self.update(0.0)

        # Log multirotor information
        self._log_multirotor_info()

    def _create_thruster_buffers(self):
        """Create thruster buffers with correct size."""
        num_instances = self.num_instances
        num_thrusters = self._count_thrusters_from_config()

        # Create thruster data tensors with correct size
        self._data.default_thruster_rps = torch.zeros(num_instances, num_thrusters, device=self.device)
        # thrust after controller and allocation is applied
        self._data.thrust_target = torch.zeros(num_instances, num_thrusters, device=self.device)
        self._data.computed_thrust = torch.zeros(num_instances, num_thrusters, device=self.device)
        self._data.applied_thrust = torch.zeros(num_instances, num_thrusters, device=self.device)

        # Combined wrench buffers
        self._thrust_target_sim = torch.zeros_like(self._data.thrust_target)  # thrust after actuator model is applied
        # wrench target for combined mode
        self._internal_wrench_target_sim = torch.zeros(num_instances, 6, device=self.device)
        # internal force/torque targets per body for combined mode
        self._internal_force_target_sim = torch.zeros(num_instances, self.num_bodies, 3, device=self.device)
        self._internal_torque_target_sim = torch.zeros(num_instances, self.num_bodies, 3, device=self.device)

        # Placeholder thruster names (will be filled during actuator creation)
        self._data.thruster_names = [f"thruster_{i}" for i in range(num_thrusters)]

    def _count_thrusters_from_config(self) -> int:
        """Count total number of thrusters from actuator configuration.

        This method parses all actuator configurations to determine the total number
        of thrusters before they are initialized. It uses the thruster name expressions
        to find matching bodies in the USD prim.

        Returns:
            Total number of thrusters across all actuator groups.

        Raises:
            ValueError: If no thrusters are found in the configuration.
        """
        total_thrusters = 0

        for actuator_name, actuator_cfg in self.cfg.actuators.items():
            if not hasattr(actuator_cfg, "thruster_names_expr"):
                continue

            # Use find_bodies to count thrusters for this actuator
            body_indices, thruster_names = self.find_bodies(actuator_cfg.thruster_names_expr)
            total_thrusters += len(body_indices)

        if total_thrusters == 0:
            raise ValueError(
                "No thrusters found in actuator configuration. "
                "Please check 'thruster_names_expr' in the provided 'MultirotorCfg.actuators' configuration."
            )

        return total_thrusters

    def _process_actuators_cfg(self):
        """Override parent method to do nothing - we handle thrusters separately."""
        # Do nothing - we handle thruster processing in _process_thruster_cfg() otherwise this
        # gives issues with joint name expressions
        pass

    def _process_cfg(self):
        """Post processing of multirotor configuration parameters."""
        # Handle root state (like parent does)
        default_root_state = (
            tuple(self.cfg.init_state.pos)
            + tuple(self.cfg.init_state.rot)
            + tuple(self.cfg.init_state.lin_vel)
            + tuple(self.cfg.init_state.ang_vel)
        )
        default_root_state = torch.tensor(default_root_state, dtype=torch.float, device=self.device)
        self._data.default_root_state = default_root_state.repeat(self.num_instances, 1)

        # Handle thruster-specific initial state
        if hasattr(self._data, "default_thruster_rps") and hasattr(self.cfg.init_state, "rps"):
            # Match against thruster names
            indices_list, _, values_list = string_utils.resolve_matching_names_values(
                self.cfg.init_state.rps, self.thruster_names
            )
            if indices_list:
                rps_values = torch.tensor(values_list, device=self.device)
                self._data.default_thruster_rps[:, indices_list] = rps_values
                self._data.thrust_target[:, indices_list] = rps_values

    def _process_thruster_cfg(self):
        """Process and apply multirotor thruster properties."""
        # create actuators
        self.actuators = dict()
        self._has_implicit_actuators = False

        # Check for mixed configurations (same as before)
        has_thrusters = False
        has_joints = False

        for actuator_name, actuator_cfg in self.cfg.actuators.items():
            if hasattr(actuator_cfg, "thruster_names_expr"):
                has_thrusters = True
            elif hasattr(actuator_cfg, "joint_names_expr"):
                has_joints = True

        if has_thrusters and has_joints:
            raise ValueError("Mixed configurations with both thrusters and regular joints are not supported.")

        if has_joints:
            raise ValueError("Regular joint actuators are not supported in Multirotor class.")

        # Store the body-to-thruster mapping
        self._thruster_body_mapping = {}

        # Track thruster names as we create actuators
        all_thruster_names = []

        for actuator_name, actuator_cfg in self.cfg.actuators.items():
            body_indices, thruster_names = self.find_bodies(actuator_cfg.thruster_names_expr)

            # Create 0-based thruster array indices starting from current count
            start_idx = len(all_thruster_names)
            thruster_array_indices = list(range(start_idx, start_idx + len(body_indices)))

            # Track all thruster names
            all_thruster_names.extend(thruster_names)

            # Store the mapping
            self._thruster_body_mapping[actuator_name] = {
                "body_indices": body_indices,
                "array_indices": thruster_array_indices,
                "thruster_names": thruster_names,
            }

            # Create thruster actuator
            actuator: Thruster = actuator_cfg.class_type(
                cfg=actuator_cfg,
                thruster_names=thruster_names,
                thruster_ids=thruster_array_indices,
                num_envs=self.num_instances,
                device=self.device,
                init_thruster_rps=self._data.default_thruster_rps[:, thruster_array_indices],
            )

            # Store actuator
            self.actuators[actuator_name] = actuator

            # Log information
            logger.info(
                f"Thruster actuator: {actuator_name} with model '{actuator_cfg.class_type.__name__}'"
                f" (thruster names: {thruster_names} [{body_indices}])."
            )

        # Update thruster names in data container
        self._data.thruster_names = all_thruster_names

        # Log summary
        logger.info(f"Initialized {len(self.actuators)} thruster actuator(s) for multirotor.")

    def _apply_actuator_model(self):
        """Processes thruster commands for the multirotor by forwarding them to the actuators.

        This internal method iterates through all thruster actuator groups and applies their
        respective actuator models to the thrust targets. The actuator models simulate realistic
        motor dynamics including:

        - Rise/fall time constants for asymmetric response
        - Thrust saturation and clipping to physical limits
        - Integration of motor dynamics over time

        The computed thrust values are stored in internal buffers for subsequent wrench computation.

        Note:
            This method updates:
            - :attr:`_thrust_target_sim`: Processed thrust values after actuator model
            - :attr:`_data.computed_thrust`: Thrust before saturation
            - :attr:`_data.applied_thrust`: Final thrust after saturation
        """

        # process thruster actions per group
        for actuator in self.actuators.values():
            if not isinstance(actuator, Thruster):
                continue

            # prepare input for actuator model based on cached data
            control_action = MultiRotorActions(
                thrusts=self._data.thrust_target[:, actuator.thruster_indices],
                thruster_indices=actuator.thruster_indices,
            )

            # compute thruster command from the actuator model
            control_action = actuator.compute(control_action)

            # update targets (these are set into the simulation)
            if control_action.thrusts is not None:
                self._thrust_target_sim[:, actuator.thruster_indices] = control_action.thrusts

            # update state of the actuator model
            self._data.computed_thrust[:, actuator.thruster_indices] = actuator.computed_thrust
            self._data.applied_thrust[:, actuator.thruster_indices] = actuator.applied_thrust

    def _apply_combined_wrench(self):
        """Apply combined wrench to the base link.

        This internal method applies the 6D wrench (computed by :meth:`_combine_thrusts`)
        to the base link of the multirotor. The wrench is applied at the center of mass
        of the base link in the local body frame.

        The forces and torques are applied through PhysX's force/torque API, which integrates
        them during the physics step to produce accelerations and velocities.
        """
        # Combine individual thrusts into a wrench vector
        self._combine_thrusts()

        self.root_physx_view.apply_forces_and_torques_at_position(
            force_data=self._internal_force_target_sim.view(-1, 3),  # Shape: (num_envs * num_bodies, 3)
            torque_data=self._internal_torque_target_sim.view(-1, 3),  # Shape: (num_envs * num_bodies, 3)
            position_data=None,  # Apply at center of mass
            indices=self._ALL_INDICES,
            is_global=False,  # Forces are in local frame
        )

    def _combine_thrusts(self):
        """Combine individual thrusts into a wrench vector.

        This internal method uses the allocation matrix to convert individual thruster
        forces into a 6D wrench vector (3D force + 3D torque) in the body frame. The
        wrench is then assigned to the base link (body index 0) for application to
        the simulation.

        The allocation matrix encodes the geometric configuration of the thrusters,
        including their positions and orientations relative to the center of mass.

        Mathematical operation:
            wrench = allocation_matrix @ thrusts
            where wrench = [Fx, Fy, Fz, Tx, Ty, Tz]^T
        """
        thrusts = self._thrust_target_sim
        self._internal_wrench_target_sim = (self.allocation_matrix @ thrusts.T).T
        # Apply forces to base link (body index 0) only
        self._internal_force_target_sim[:, 0, :] = self._internal_wrench_target_sim[:, :3]
        self._internal_torque_target_sim[:, 0, :] = self._internal_wrench_target_sim[:, 3:]

    def _validate_cfg(self):
        """Validate the multirotor configuration after processing.

        Note:
            This function should be called only after the configuration has been processed and the buffers have been
            created. Otherwise, some settings that are altered during processing may not be validated.
        """
        # Only validate if actuators have been created
        if hasattr(self, "actuators") and self.actuators:
            # Validate thruster-specific configuration
            for actuator_name in self.actuators:
                if isinstance(self.actuators[actuator_name], Thruster):
                    initial_thrust = self.actuators[actuator_name].curr_thrust
                    # check that the initial thrust is within the limits
                    thrust_limits = self.actuators[actuator_name].cfg.thrust_range
                    if torch.any(initial_thrust < thrust_limits[0]) or torch.any(initial_thrust > thrust_limits[1]):
                        raise ValueError(
                            f"Initial thrust for actuator '{actuator_name}' is out of bounds: "
                            f"{initial_thrust} not in {thrust_limits}"
                        )

    def _log_multirotor_info(self):
        """Log multirotor-specific information."""
        logger.info(f"Multirotor initialized with {self.num_thrusters} thrusters")
        logger.info(f"Thruster names: {self.thruster_names}")
        logger.info(f"Thruster force direction: {self.cfg.thruster_force_direction}")
