from collections.abc import Callable
from enum import Enum
from itertools import pairwise
from typing import Annotated, Any, ClassVar

import numpy as np
from numpy.typing import NDArray
from pydantic import BeforeValidator

from characterization.utils.scenario_types import AgentType

SMALL_EPS = 1e-6
BIG_EPS = 1e6
SUPPORTED_SCENARIO_TYPES = ["gt", "ho"]
SUPPORTED_CRITERIA = ["critical", "average"]
MIN_VALID_POINTS = 2
MAX_DECELERATION = 10.0  # m/s^2,
STATIONARY_SPEED_THRESHOLD = 0.5  # m/s
SUPPORTED_SCORERS = ["individual", "interaction", "safeshift"]


def categorize_from_thresholds(value: float, threshold_values: list[float]) -> int:
    """Categorizes a value based on provided ranges.

    Args:
        value (float): The value to categorize.
        threshold_values (list[float]): A list of threshold values defining the ranges.

    Returns:
        int: The category index (1, 2, ..., n+1) based on the ranges.
    """
    num_thresholds = len(threshold_values)
    assert num_thresholds >= 1, "At least one range must be provided."

    # If there is only one category, return 1 or 2 based on the value
    if num_thresholds < 2:  # noqa: PLR2004
        return 1 if value <= threshold_values[0] else 2

    # If value is below the lowest range, return 1
    if value <= threshold_values[0]:
        return 1

    # Categorize based on ranges, starting from category 2
    for category, (lower_bound, upper_bound) in enumerate(pairwise(threshold_values)):
        if lower_bound < value <= upper_bound:
            return category + 2

    # If value is above the highest range
    return num_thresholds + 1


def mph_to_ms(mph: float) -> float:
    """Converts miles per hour (mph) to meters per second (m/s).

    Args:
        mph (float): Speed in miles per hour.

    Returns:
        float: Speed in meters per second.
    """
    return mph * 0.44704  # 1 mph = 0.44704 m/s


# Validator factory
def validate_array(
    expected_dtype: Any,  # noqa: ANN401
    expected_ndim: int,
) -> Callable[[Any], NDArray[Any]]:
    """Factory function to create a validator for numpy arrays with specific dtype and ndim."""

    def _validator(
        v: Any,  # noqa: ANN401
    ) -> NDArray[Any]:
        if not isinstance(v, np.ndarray):
            error_message = f"Expected a numpy.ndarray, got {type(v)}"
            raise TypeError(error_message)
        if v.dtype != expected_dtype:
            error_message = f"Expected dtype {expected_dtype}, got {v.dtype}"
            raise TypeError(error_message)
        if v.ndim != expected_ndim:
            error_message = f"Expected {expected_ndim}D array, got {v.ndim}D"
            raise ValueError(error_message)
        return v

    return _validator


# Reusable types
BooleanNDArray1D = Annotated[NDArray[np.bool_], BeforeValidator(validate_array(np.bool_, 1))]
BooleanNDArray2D = Annotated[NDArray[np.bool_], BeforeValidator(validate_array(np.bool_, 2))]
BooleanNDArray3D = Annotated[NDArray[np.bool_], BeforeValidator(validate_array(np.bool_, 3))]
Float64NDArray3D = Annotated[NDArray[np.float64], BeforeValidator(validate_array(np.float64, 3))]
Float32NDArray4D = Annotated[NDArray[np.float32], BeforeValidator(validate_array(np.float32, 4))]
Float32NDArray3D = Annotated[NDArray[np.float32], BeforeValidator(validate_array(np.float32, 3))]
Float32NDArray2D = Annotated[NDArray[np.float32], BeforeValidator(validate_array(np.float32, 2))]
Float32NDArray1D = Annotated[NDArray[np.float32], BeforeValidator(validate_array(np.float32, 1))]
Int32NDArray1D = Annotated[NDArray[np.int32], BeforeValidator(validate_array(np.int32, 1))]
Int32NDArray2D = Annotated[NDArray[np.int32], BeforeValidator(validate_array(np.int32, 2))]
Int64NDArray2D = Annotated[NDArray[np.int64], BeforeValidator(validate_array(np.int64, 2))]


class FeatureType(Enum):
    """Enumeration for feature types."""

    CONTINUOUS = "continuous"
    CATEGORICAL = "categorical"


class InteractionStatus(Enum):
    """Enumeration for interaction status."""

    UNKNOWN = -1
    COMPUTED_OK = 0
    PARTIAL_INVALID_HEADING = 1
    MASK_NOT_VALID = 2
    DISTANCE_TOO_FAR = 3
    STATIONARY = 4


class ReturnCriterion(Enum):
    """Enumeration for return criteria."""

    CRITICAL = 0
    AVERAGE = 1
    UNSET = -1


class AgentTrajectoryMasker:
    """Masks for indexing trajectory data from the reformatted by the dataloader classes.

    The class expects an input of type (N, T, D=10) or (T, D=10) where N is the number of agents, T is the number of
    timesteps and D is the number of features per trajectory point, organized as follows:
        idx 0 to 2: the agent's (x, y, z) center coordinates.
        idx 3 to 5: the agent's length, width and height in meters.
        idx 6: the agent's angle (heading) of the forward direction in radians
        idx 7 to 8: the agent's (x, y) velocity in meters/second
        idx 9: a flag indicating if the information is valid
    """

    # Agent position masks
    _TRAJECTORY_XYZ_POS: ClassVar[list[bool]] = [True, True, True, False, False, False, False, False, False, False]
    _TRAJECTORY_XY_POS: ClassVar[list[bool]] = [True, True, False, False, False, False, False, False, False, False]

    # Agent dimensions masks
    _TRAJECTORY_DIMS: ClassVar[list[bool]] = [False, False, False, True, True, True, False, False, False, False]
    _TRAJECTORY_LENGTHS: ClassVar[list[bool]] = [False, False, False, True, False, False, False, False, False, False]
    _TRAJECTORY_WIDTHS: ClassVar[list[bool]] = [False, False, False, False, True, False, False, False, False, False]
    _TRAJECTORY_HEIGHTS: ClassVar[list[bool]] = [False, False, False, False, False, True, False, False, False, False]

    # Agent heading mask
    _TRAJECTORY_HEADING: ClassVar[list[bool]] = [False, False, False, False, False, False, True, False, False, False]

    # Agent velocity masks
    _TRAJECTORY_XY_VEL: ClassVar[list[bool]] = [False, False, False, False, False, False, False, True, True, False]
    _TRAJECTORY_X_VEL: ClassVar[list[bool]] = [False, False, False, False, False, False, False, True, False, False]
    _TRAJECTORY_Y_VEL: ClassVar[list[bool]] = [False, False, False, False, False, False, False, False, True, False]

    # Agent state, all features except valid mask
    _TRAJECTORY_STATE: ClassVar[list[bool]] = [True, True, True, True, True, True, True, True, True, False]

    # Agent valid mask
    _TRAJECTORY_VALID: ClassVar[list[bool]] = [False, False, False, False, False, False, False, False, False, True]

    _agent_trajectory: NDArray[np.float32]

    def __init__(self, trajectory: NDArray[np.float32]) -> None:
        """Initializes the AgentTrajectoryMasker with trajectory data.

        Args:
            trajectory (NDArray[np.float32]): The trajectory data of shape (N, T, D=10) or (T, D=10).
        """
        self._agent_trajectory = trajectory

    # Mask accessors
    @property
    def xyz_pos_mask(self) -> list[bool]:
        """Mask for the (x, y, z) position feature."""
        return self._TRAJECTORY_XYZ_POS

    @property
    def xy_pos_mask(self) -> list[bool]:
        """Mask for the (x, y) position feature."""
        return self._TRAJECTORY_XY_POS

    @property
    def xy_vel_mask(self) -> list[bool]:
        """Mask for the (x, y) velocity feature."""
        return self._TRAJECTORY_XY_VEL

    @property
    def heading_mask(self) -> list[bool]:
        """Mask for the heading feature."""
        return self._TRAJECTORY_HEADING

    # Trajectory accessors
    @property
    def agent_trajectories(self) -> NDArray[np.float32]:
        """Returns the full agent trajectory data."""
        return self._agent_trajectory

    @property
    def agent_dims(self) -> NDArray[np.float32]:
        """Returns the agents dimensions: length, width, height."""
        return self._agent_trajectory[..., self._TRAJECTORY_DIMS]

    @property
    def agent_lengths(self) -> NDArray[np.float32]:
        """Returns the length."""
        return self._agent_trajectory[..., self._TRAJECTORY_LENGTHS]

    @property
    def agent_widths(self) -> NDArray[np.float32]:
        """Returns the width."""
        return self._agent_trajectory[..., self._TRAJECTORY_WIDTHS]

    @property
    def agent_heights(self) -> NDArray[np.float32]:
        """Returns the height."""
        return self._agent_trajectory[..., self._TRAJECTORY_HEIGHTS]

    @property
    def agent_headings(self) -> NDArray[np.float32]:
        """Returns the heading."""
        return self._agent_trajectory[..., self._TRAJECTORY_HEADING]

    @property
    def agent_xyz_pos(self) -> NDArray[np.float32]:
        """Returns the (x, y, z) position."""
        return self._agent_trajectory[..., self._TRAJECTORY_XYZ_POS]

    @property
    def agent_xy_pos(self) -> NDArray[np.float32]:
        """Returns the (x, y) position."""
        return self._agent_trajectory[..., self._TRAJECTORY_XY_POS]

    @property
    def agent_xy_vel(self) -> NDArray[np.float32]:
        """Returns the (x, y) velocity."""
        return self._agent_trajectory[..., self._TRAJECTORY_XY_VEL]

    @property
    def agent_valid(self) -> NDArray[np.float32]:
        """Returns the valid mask."""
        valid = self._agent_trajectory[..., self._TRAJECTORY_VALID]
        return np.nan_to_num(valid, nan=0.0)

    @property
    def agent_state(self) -> NDArray[np.float32]:
        """Returns all features except the valid mask."""
        return self._agent_trajectory[..., self._TRAJECTORY_STATE]


class LaneMasker:
    """Masks for indexing lane data from the reformatted by the dataloader classes.

    The class expects an input of shape (N, L, T, D=6) or (L, T, D=6) where N is the number of agents, L is the number
    of lanes, and T is the number of timesteps. D is the number of features per lane point, organized as follows:
    timesteps and D is the number of features per lane point, organized as follows:
        idx 0: closest lane distance to the agent in meters.
        idx 1: lane point index of the closest lane point to the agent.
        idx 2 to 4: the lane point's (x, y, z) coordinates.
        idx 5: lane index.
    """

    # Lane Distances
    _LANE_DISTS: ClassVar[list[bool]] = [True, False, False, False, False, False]

    # Lane Point Index
    _LANE_POINT_IDX: ClassVar[list[bool]] = [False, True, False, False, False, False]

    # Lane Point (x, y, z) position masks
    _LANE_POINT_XYZ_POS: ClassVar[list[bool]] = [False, False, True, True, True, False]
    _LANE_POINT_XY_POS: ClassVar[list[bool]] = [False, False, True, True, False, False]

    # Lane Index
    _LANE_IDX: ClassVar[list[bool]] = [False, False, False, False, False, True]

    # Lane and distance
    _LANE_DIST_AND_IDX: ClassVar[list[bool]] = [True, False, False, False, False, True]

    _lane_to_agent_metadata: NDArray[np.float32]

    def __init__(self, lane_to_agent_metadata: NDArray[np.float32]) -> None:
        """Initializes the LaneMasker with lane to agent metadata.

        Args:
            lane_to_agent_metadata (NDArray[np.float32]): The lane to agent metadata of shape (N, T, D=6) or (T, D=6).
        """
        self._lane_to_agent_metadata = lane_to_agent_metadata

    # Lane metadata accessors
    @property
    def lane_to_agent_metadata(self) -> NDArray[np.float32]:
        """Returns the lane to agent metadata."""
        return self._lane_to_agent_metadata

    @property
    def lane_dists(self) -> NDArray[np.float32]:
        """Returns the closest lane distances to the agent."""
        return self._lane_to_agent_metadata[..., self._LANE_DISTS]

    @property
    def lane_point_idx(self) -> NDArray[np.float32]:
        """Returns the lane point indices of the closest lane points to the agent."""
        return self._lane_to_agent_metadata[..., self._LANE_POINT_IDX]

    @property
    def lane_point_xyz_pos(self) -> NDArray[np.float32]:
        """Returns the (x, y, z) position of the closest lane points to the agent."""
        return self._lane_to_agent_metadata[..., self._LANE_POINT_XYZ_POS]

    @property
    def lane_point_xy_pos(self) -> NDArray[np.float32]:
        """Returns the (x, y) position of the closest lane points to the agent."""
        return self._lane_to_agent_metadata[..., self._LANE_POINT_XY_POS]

    @property
    def lane_idx(self) -> NDArray[int]:
        """Returns the lane indices of the closest lanes to the agent."""
        return self._lane_to_agent_metadata[..., self._LANE_IDX].astype(int)

    @property
    def lane_dist_and_idx(self) -> NDArray[np.float32]:
        """Returns the closest lane distances and lane indices to the agent."""
        return self._lane_to_agent_metadata[..., self._LANE_DIST_AND_IDX]


class InteractionAgent:
    """Class representing an agent for interaction feature computation."""

    def __init__(self) -> None:
        """Initializes an InteractionAgent and resets all attributes."""
        self.reset()

    @property
    def position(self) -> NDArray[np.float32]:
        """NDArray[np.float32]: The positions of the agent over time (shape: [T, 2])."""
        return self._position

    @position.setter
    def position(self, value: NDArray[np.float32]) -> None:
        """Sets the positions of the agent.

        Args:
            value (NDArray[np.float32]): The positions of the agent over time (shape: [T, 2]).
        """
        self._position = np.asarray(value, dtype=np.float32)

    @property
    def speed(self) -> NDArray[np.float32]:
        """NDArray[np.float32]: The velocities of the agent over time (shape: [T,])."""
        return self._speed

    @speed.setter
    def speed(self, value: NDArray[np.float32]) -> None:
        """Sets the velocities of the agent.

        Args:
            value (NDArray[np.float32]): The velocities of the agent over time (shape: [T,]).
        """
        self._speed = np.asarray(value, dtype=np.float32)

    @property
    def heading(self) -> NDArray[np.float32]:
        """NDArray[np.float32]: The headings of the agent over time (shape: [T,])."""
        return self._heading

    @heading.setter
    def heading(self, value: NDArray[np.float32]) -> None:
        """Sets the headings of the agent.

        Args:
            value (NDArray[np.float32]): The headings of the agent over time (shape: [T,]).
        """
        self._heading = np.asarray(value, dtype=np.float32)

    @property
    def length(self) -> NDArray[np.float32]:
        """NDArray[np.float32]: The lengths of the agent over time (shape: [T,])."""
        return self._length

    @length.setter
    def length(self, value: NDArray[np.float32]) -> None:
        """Sets the lengths of the agent.

        Args:
            value (NDArray[np.float32]): The lengths of the agent over time (shape: [T,]).
        """
        self._length = np.asarray(value, dtype=np.float32)

    @property
    def width(self) -> NDArray[np.float32]:
        """NDArray[np.float32]: The widths of the agent over time (shape: [T,])."""
        return self._width

    @width.setter
    def width(self, value: NDArray[np.float32]) -> None:
        """Sets the widths of the agent.

        Args:
            value (NDArray[np.float32]): The widths of the agent over time (shape: [T,]).
        """
        self._width = np.asarray(value, dtype=np.float32)

    @property
    def height(self) -> NDArray[np.float32]:
        """NDArray[np.float32]: The heights of the agent over time (shape: [T,])."""
        return self._height

    @height.setter
    def height(self, value: NDArray[np.float32]) -> None:
        """Sets the heights of the agent.

        Args:
            value (NDArray[np.float32]): The heights of the agent over time (shape: [T,]).
        """
        self._height = np.asarray(value, dtype=np.float32)

    @property
    def agent_type(self) -> AgentType:
        """str: The type of the agent."""
        return self._agent_type

    @agent_type.setter
    def agent_type(self, value: AgentType) -> None:
        """Sets the type of the agent.

        Args:
            value (str): The type of the agent.
        """
        self._agent_type = value

    @property
    def is_stationary(self) -> bool:
        """Bool: Whether the agent is stationary (True/False)."""
        self._is_stationary = self.speed.mean() < self._stationary_speed
        return self._is_stationary

    @property
    def stationary_speed(self) -> float:
        """float: The speed threshold below which the agent is considered stationary."""
        return self._stationary_speed

    @stationary_speed.setter
    def stationary_speed(self, value: float) -> None:
        """Sets the stationary speed threshold.

        Args:
            value (float): The speed threshold below which the agent is considered stationary.
        """
        self._stationary_speed = value

    @property
    def in_conflict_point(self) -> bool:
        """bool: Whether the agent is in a conflict point."""
        self._in_conflict_point = np.any(
            self._dists_to_conflict <= self._agent_to_conflict_point_max_distance
        ).__bool__()
        return self._in_conflict_point

    @property
    def agent_to_conflict_point_max_distance(self) -> float:
        """float: The maximum distance to a conflict point."""
        return self._agent_to_conflict_point_max_distance

    @agent_to_conflict_point_max_distance.setter
    def agent_to_conflict_point_max_distance(self, value: float) -> None:
        """Sets the maximum distance to a conflict point.

        Args:
            value (float): The maximum distance to a conflict point.
        """
        self._agent_to_conflict_point_max_distance = value

    @property
    def dists_to_conflict(self) -> NDArray[np.float32]:
        """NDArray[np.float32]: The distances to conflict points (shape: [T,])."""
        return self._dists_to_conflict

    @dists_to_conflict.setter
    def dists_to_conflict(self, value: NDArray[np.float32] | None) -> None:
        """Sets the distances to conflict points.

        Args:
            value (NDArray[np.float32] | None): The distances to conflict points (shape: [T,]).
        """
        self._dists_to_conflict = np.asarray(value, dtype=np.float32)

    @property
    def lane(self) -> NDArray[np.float32] | None:
        """NDArray[np.float32] or None: The lane of the agent, if available."""
        return self._lane

    @lane.setter
    def lane(self, value: NDArray[np.float32] | None) -> None:
        """Sets the lane of the agent.

        Args:
            value (NDArray[np.float32] | None): The lane of the agent, if available.
        """
        if value is not None:
            self._lane = np.asarray(value, dtype=np.float32)
        else:
            self._lane = None

    def reset(self) -> None:
        """Resets all agent attributes to their default values."""
        self._position = np.empty((0, 2), dtype=np.float32)
        self._speed = np.empty((0,), dtype=np.float32)
        self._heading = np.empty((0,), dtype=np.float32)
        self._dists_to_conflict = np.empty((0,), dtype=np.float32)
        self._stationary_speed = 0.1  # Default stationary speed threshold
        self._agent_to_conflict_point_max_distance = 0.5  # Default max distance to conflict point
        self._lane = np.empty((0,), dtype=np.float32)
        self._length = np.empty((0,), dtype=np.float32)
        self._width = np.empty((0,), dtype=np.float32)
        self._height = np.empty((0,), dtype=np.float32)
        self._agent_type = AgentType.TYPE_UNSET


class TrajectoryType(Enum):
    """Trajectory Types for WOMD."""

    TYPE_UNSET = -1
    TYPE_STATIONARY = 0
    TYPE_STRAIGHT = 1
    TYPE_STRAIGHT_RIGHT = 2
    TYPE_STRAIGHT_LEFT = 3
    TYPE_RIGHT_U_TURN = 4
    TYPE_RIGHT_TURN = 5
    TYPE_LEFT_U_TURN = 6
    TYPE_LEFT_TURN = 7


# Weights for different trajectory types are loosely set based on Figure 3 (a) of https://arxiv.org/pdf/2403.15098
# Weight per class is set as (100% - class frequency %) * 0.10
TRAJECTORY_TYPE_WEIGHTS = {
    TrajectoryType.TYPE_UNSET: 0.0,
    # Stationary agents correspond to less than 10% of the data
    TrajectoryType.TYPE_STATIONARY: 9.0,
    # Straight-moving agents correspond to ~50% of the data.
    TrajectoryType.TYPE_STRAIGHT: 5.0,
    # Straight-right agents correspond less than ~10% of the data.
    TrajectoryType.TYPE_STRAIGHT_RIGHT: 9.0,
    # Straight-left agents correspond to ~10% of the data.
    TrajectoryType.TYPE_STRAIGHT_LEFT: 9.0,
    # Right-turn agents correspond to less than 20% of the data.
    TrajectoryType.TYPE_RIGHT_TURN: 8.0,
    # Left-turn agents correspond to less than 20% of the data.
    TrajectoryType.TYPE_LEFT_TURN: 8.0,
    # Right-U-turn agents correspond to less than 10% of the data.
    TrajectoryType.TYPE_RIGHT_U_TURN: 9.0,
    # Left-U-turn agents correspond to less than 10% of the data.
    TrajectoryType.TYPE_LEFT_U_TURN: 9.0,
}
