# Copyright (c) 2022-2023, The ORBIT Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

from __future__ import annotations

import torch
from dataclasses import dataclass


@dataclass
class ContactSensorData:
    """Data container for the contact reporting sensor."""

    pos_w: torch.Tensor | None = None
    """Position of the sensor origin in world frame.

    Shape is (N, 3), where ``N`` is the number of sensors.

    Note:
        If the :attr:`ContactSensorCfg.track_pose` is False, then this qunatity is None.
    """

    quat_w: torch.Tensor | None = None
    """Orientation of the sensor origin in quaternion ``(w, x, y, z)`` in world frame.

    Shape is (N, 4), where ``N`` is the number of sensors.

    Note:
        If the :attr:`ContactSensorCfg.track_pose` is False, then this qunatity is None.
    """

    net_forces_w: torch.Tensor = None
    """The net contact forces in world frame.

    Shape is (N, B, 3), where ``N`` is the number of sensors and ``B`` is the number of bodies in each sensor.
    """

    net_forces_w_history: torch.Tensor = None
    """The net contact forces in world frame.

    Shape is (N, T, B, 3), where ``N`` is the number of sensors, ``T`` is the configured history length
    and ``B`` is the number of bodies in each sensor.

    In the history dimension, the first index is the most recent and the last index is the oldest.
    """

    force_matrix_w: torch.Tensor | None = None
    """The contact forces filtered between the sensor bodies and filtered bodies in world frame.

    Shape is (N, B, S, M, 3), where ``N`` is the number of sensors, ``B`` is number of bodies in each sensor,
    ``S`` is number of shapes per body and ``M`` is the number of filtered bodies.

    Note:
        If the :attr:`ContactSensorCfg.filter_prim_paths_expr` is empty, then this quantity is None.
    """

    last_air_time: torch.Tensor | None = None
    """Time spent (in s) in the air before the last contact.

    Shape is (N,), where ``N`` is the number of sensors.

    Note:
        If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None.
    """

    current_air_time: torch.Tensor | None = None
    """Time spent (in s) in the air since the last contact.

    Shape is (N,), where ``N`` is the number of sensors.

    Note:
        If the :attr:`ContactSensorCfg.track_air_time` is False, then this quantity is None.
    """
