"""
Compute matrix multiplication of two matrices represented as nested lists.

The module provides a single function :func:`matmul` that performs the classic
O(n³) algorithm without any external dependencies.

Mathematics
-----------

If ``A`` has shape ``(m, k)`` and ``B`` has shape ``(k, n)``, the result
``C`` has shape ``(m, n)`` where::

    C[i][j] = sum_{t=0}^{k-1} A[i][t] * B[t][j]

This implementation follows that definition directly.

Validations
------------

* Both matrices must be non‑empty and **rectangular** – every row of a
  matrix must have the same length.
* The number of columns of ``matrix_a`` must equal the number of rows of
  ``matrix_b``.

If any of these conditions are violated a :class:`ValueError` is raised
with a descriptive message.
"""

from __future__ import annotations

import math

__all__ = ["matmul", "rotation"]


def matmul(
    matrix_a: list[list[float]], matrix_b: list[list[float]]
) -> list[list[float]]:
    """Return the matrix product of ``matrix_a`` and ``matrix_b``.

    Parameters
    ----------
    matrix_a:
        Left‑hand matrix with shape ``(m, k)``.
    matrix_b:
        Right‑hand matrix with shape ``(k, n)``.
    Returns
    -------
    List[List[float]]
        The resulting matrix with shape ``(m, n)``.
    Raises
    ------
    ValueError
        If either matrix is empty, not rectangular, or if the inner
        dimensions (columns of ``matrix_a`` and rows of ``matrix_b``)
        mismatch.
    """
    # --------- Validation of input matrices ----------
    # Check that the matrices are not empty
    if not matrix_a or not matrix_b:
        raise ValueError("Both matrices must be non‑empty")
    # Determine row and column counts for each matrix
    a_rows, a_cols = len(matrix_a), len(matrix_a[0])
    b_rows, b_cols = len(matrix_b), len(matrix_b[0])
    # Verify all rows have consistent length (rectangularity)
    if not all(len(row) == a_cols for row in matrix_a):
        raise ValueError("matrix_a is not rectangular")
    if not all(len(row) == b_cols for row in matrix_b):
        raise ValueError("matrix_b is not rectangular")
    # Verify inner dimensions match for multiplication
    if a_cols != b_rows:
        raise ValueError(
            f"Cannot multiply: matrix_a's columns ({a_cols}) "
            f"do not match matrix_b's rows ({b_rows})"
        )
    # --------- Matrix multiplication -----------
    result: list[list[float]] = []
    for i in range(a_rows):
        row_result: list[float] = []
        for j in range(b_cols):
            # Compute dot product of i‑th row of A and j‑th column of B
            dot = 0.0
            for k in range(a_cols):
                dot += matrix_a[i][k] * matrix_b[k][j]
            row_result.append(dot)
        result.append(row_result)
    return result


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

    Parameters
    ----------
    matrix
        A non‑empty list of lists of floats where each inner list has
        the same length.

    Returns
    -------
    List[List[float]]
        The transposed matrix, where rows become columns and
        columns become rows.
    Raises
    ------
    ValueError
        If the matrix is empty or not rectangular.
    """
    if not matrix:
        raise ValueError("Matrix must be non‑empty")
    rows, cols = len(matrix), len(matrix[0])
    if not all(len(row) == cols for row in matrix):
        raise ValueError("Matrix must be rectangular")
    # Transpose by swapping indices
    return [[matrix[r][c] for r in range(rows)] for c in range(cols)]


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

    The angle is specified in radians.

    Parameters
    ----------
    angle : float
        Rotation angle in radians.

    Returns
    -------
    list[list[float]]
        A 2x2 matrix
        ``[[cos(angle), -sin(angle)],
        [sin(angle),  cos(angle)]]``
    """
    cos_theta = math.cos(angle)
    sin_theta = math.sin(angle)
    return [
        [cos_theta, -sin_theta],
        [sin_theta, cos_theta],
    ]


if __name__ == "__main__":
    # Example 1: 2x3 * 3x2 = 2x2
    A = [[1, 2, 3], [4, 5, 6]]
    B = [[7, 8], [9, 10], [11, 12]]
    expected_1 = [[58, 64], [139, 154]]
    assert matmul(A, B) == expected_1, f"Test 1 failed: got {matmul(A, B)}"

    # Example 2: 2x2 * 2x2
    C = [[1, 2], [3, 4]]
    D = [[5, 6], [7, 8]]
    expected_2 = [[19, 22], [43, 50]]
    assert matmul(C, D) == expected_2, f"Test 2 failed: got {matmul(C, D)}"
    assert matmul(C, D) == expected_2, f"Test 2 failed: got {matmul(C, D)}"
    # Example 3: Transpose a 2x2 matrix
    M = [[1, 2], [3, 4]]
    expected_transpose_3 = [[1, 3], [2, 4]]
    assert transpose(M) == expected_transpose_3, (
        f"Transposition failed, got {transpose(M)}"
    )

    # Example 4: Transpose a 3x1 matrix
    N = [[5], [6], [7]]
    expected_transpose_4 = [[5, 6, 7]]
    assert transpose(N) == expected_transpose_4, (
        f"Transposition failed, got {transpose(N)}"
    )

    # Rotation tests
    angle_zero = 0.0
    mat0 = rotation(angle_zero)
    expected0 = [[1.0, 0.0], [0.0, 1.0]]
    for i in range(2):
        for j in range(2):
            assert math.isclose(mat0[i][j], expected0[i][j]), (
                f"rotation(0) mismatch at {i},{j}"
            )

    angle90 = math.pi / 2
    mat90 = rotation(angle90)
    # Expected: [[0, -1], [1, 0]] with float precision
    assert math.isclose(mat90[0][0], 0.0, rel_tol=1e-9, abs_tol=1e-9), (
        "rotation(pi/2) [0][0] mismatch"
    )
    assert math.isclose(mat90[0][1], -1.0, rel_tol=1e-9, abs_tol=1e-9), (
        "rotation(pi/2) [0][1] mismatch"
    )
    assert math.isclose(mat90[1][0], 1.0, rel_tol=1e-9, abs_tol=1e-9), (
        "rotation(pi/2) [1][0] mismatch"
    )
    assert math.isclose(mat90[1][1], 0.0, rel_tol=1e-9, abs_tol=1e-9), (
        "rotation(pi/2) [1][1] mismatch"
    )

    print("All tests passed.")