"""2D and 3D vector and matrix rotation methods."""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    from collections.abc import Sequence

    from numpy.typing import NDArray


def _matrix_3d(theta: float, axis: int) -> NDArray[np.floating]:
    """Return the 3D rotation matrix theta (radians) counterclockwise.

    EG: in an xyz frame, axis 0 would rotate the x axis counterclockwise by theta
    radians, rotating the yz-plane.

    Args:
        theta(float): Radians couterclockwise to rotate.
        axis(int): Which axis to rotate around of 0,1,2

    Returns:
        3x3 rotation matrix
    """
    if not (0 <= axis <= 2):  # noqa: PLR2004
        raise ValueError("Only valid for axis 0, 1, or 2: 3x3 matrix rotation")
    c = np.cos(theta)
    s = np.sin(theta)
    if axis == 0:
        return np.array([[1, 0, 0], [0, c, -s], [0, s, c]])
    if axis == 1:
        return np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]])
    return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])


def _matrix_2d(theta: float) -> NDArray[np.floating]:
    """Return the 2D rotation matrix for theta(radians) counterclockwise.

    Args:
        theta(float): Radians couterclockwise to rotate.

    Returns:
        2x2 rotation matrix
    """
    c = np.cos(theta)
    s = np.sin(theta)
    return np.array(((c, -s), (s, c)))


def _perform_rotation(
    rot: NDArray[np.floating],
    mat: NDArray[np.floating] | Sequence,
) -> NDArray[np.floating]:
    """Perform the rotation of the matrix or vector and return."""
    arr = rot @ mat
    if arr.ndim == 1:
        return arr
    return arr @ rot.T


def rotate_ccw(
    theta: float,
    mat: NDArray[np.floating] | Sequence[float | int] | Sequence[Sequence[float | int]],
    axis: int = 2,
) -> NDArray[np.floating]:
    """Rotate the matrix counterclockwise by theta radians.

    Rotates 2x2 or 3x3 matrices. If 3x3, axis is required to be one of 0, 1, or 2.

    Args:
        theta(float): Radians to rotate counterclockwise
        mat(np.ndarray): 2x2 or 3x3 matrix to rotate
        axis(int,optional): Required for 3x3, must be axis 0, 1, or 2. Defaults to 2

    Returns:
        np.ndarray: The rotate matrix
    """
    mat = np.asarray(mat)
    if mat.shape in ((2, 2), (1, 2), (2,)):
        r = _matrix_2d(theta)
        if axis != 2:  # noqa: PLR2004
            raise ValueError("Axis should not be defined as other than 2 or default")
    elif mat.shape in ((3, 3), (1, 3), (3,)):
        r = _matrix_3d(theta, axis)
    else:
        raise ValueError(
            "Only valid for size 2 or 3 square matrices or column vectors:"
            f" Got {mat.shape}",
        )
    return _perform_rotation(r, mat)


def rotate_cw(
    theta: float,
    mat: NDArray[np.floating] | Sequence[float | int] | Sequence[Sequence[float | int]],
    axis: int = 2,
) -> NDArray[np.floating]:
    """Rotate the matrix clockwise by theta radians.

    Rotates 2x2 or 3x3 matrices. If 3x3, axis is required to be one of 0, 1, or 2.

    Args:
        theta(float): Radians to rotate clockwise
        mat(np.ndarray): 2x2 or 3x3 matrix to rotate
        axis(int,optional): Required for 3x3, must be axis 0, 1, or 2. Defaults to 2

    Returns:
        np.ndarray: The rotate matrix
    """
    return rotate_ccw(-theta, mat, axis)
