import math

from tabulate import tabulate


def matmul(a: list[list[float]], b: list[list[float]]) -> list[list[float]]:
    """Multiply two non-empty rectangular matrices.

    Args:
        a: The left matrix.
        b: The right matrix.

    Returns:
        The matrix product a * b.

    Raises:
        ValueError: If either matrix is empty or ragged, or if their
            dimensions are incompatible.
    """
    _validate_matrix(a, "a")
    _validate_matrix(b, "b")

    if len(a[0]) != len(b):
        raise ValueError(
            "Matrix dimensions are incompatible: the number of columns in a "
            "must equal the number of rows in b."
        )

    return [
        [
            sum(a_row[index] * b[index][column] for index in range(len(b)))
            for column in range(len(b[0]))
        ]
        for a_row in a
    ]


def transpose(matrix: list[list[float]]) -> list[list[float]]:
    """Return the transpose of a non-empty rectangular matrix.

    Args:
        matrix: A matrix represented as rows of numbers.

    Returns:
        A new matrix whose rows are the columns of matrix.

    Raises:
        ValueError: If the matrix is empty or its rows have different lengths.
    """
    _validate_matrix(matrix, "matrix")
    return [list(column) for column in zip(*matrix)]


def rot_2D(angle: float) -> list[list[float]]:
    """Return the 2D rotation matrix for a given angle in radians.

    Args:
        angle: The rotation angle in radians.

    Returns:
        The rotation matrix for rotating points counterclockwise.
    """
    return [[math.cos(angle), -math.sin(angle)], [math.sin(angle), math.cos(angle)]]


def rot_3D(angle: float, axis: str) -> list[list[float]]:
    """Return a 3D rotation matrix around the x, y, or z axis.

    With ``c = cos(angle)`` and ``s = sin(angle)``, the matrices are::

        x: [[1, 0, 0], [0, c, -s], [0, s, c]]
        y: [[c, 0, s], [0, 1, 0], [-s, 0, c]]
        z: [[c, -s, 0], [s, c, 0], [0, 0, 1]]

    Args:
        angle: The rotation angle in radians.
        axis: The axis to rotate around: ``"x"``, ``"y"``, or ``"z"``.

    Returns:
        The 3D rotation matrix for the selected axis.

    Raises:
        ValueError: If axis is not ``"x"``, ``"y"`, or ``"z"``.
    """
    cosine = math.cos(angle)
    sine = math.sin(angle)

    if axis == "x":
        return [[1.0, 0.0, 0.0], [0.0, cosine, -sine], [0.0, sine, cosine]]
    if axis == "y":
        return [[cosine, 0.0, sine], [0.0, 1.0, 0.0], [-sine, 0.0, cosine]]
    if axis == "z":
        return [[cosine, -sine, 0.0], [sine, cosine, 0.0], [0.0, 0.0, 1.0]]

    raise ValueError("axis must be 'x', 'y', or 'z'")


def _validate_matrix(matrix: list[list[float]], name: str) -> None:
    """Raise ValueError unless matrix is non-empty and rectangular."""
    if not matrix:
        raise ValueError(f"Matrix {name} must be non-empty.")

    column_count = len(matrix[0])

    if column_count == 0:
        raise ValueError(f"Matrix {name} must have at least one column.")

    if any(len(row) != column_count for row in matrix):
        raise ValueError(f"Matrix {name} must have rows of equal length.")


if __name__ == "__main__":
    first = [[1, 2, 3], [4, 5, 6]]
    second = [[7, 8], [9, 10], [11.0, 12.0]]

    assert matmul(first, second) == [[58.0, 64.0], [139.0, 154.0]]
    print("Matrix product:", matmul(first, second))

    try:
        matmul([[1, 2]], [[1, 2]])
    except ValueError as error:
        print("Expected error:", error)
    else:
        raise AssertionError("Expected a ValueError for incompatible dimensions.")

    assert transpose(first) == [[1, 4], [2, 5], [3, 6]]
    print("Transpose:", transpose(first))

    print(tabulate(rot_2D(1.0)))