"""Lesson 2 lab: vectors, matrices, broadcasting, rank, and solving Ax=b.

Run:
    python 02-vectors-matrices-operations.py

Requires:
    numpy
"""

from __future__ import annotations

import numpy as np


def manual_matmul(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """Educational matrix multiplication with an explicit shape contract."""
    if a.ndim != 2 or b.ndim != 2 or a.shape[1] != b.shape[0]:
        raise ValueError(f"Illegal matrix product: {a.shape} @ {b.shape}")
    result = np.zeros((a.shape[0], b.shape[1]), dtype=np.result_type(a, b))
    for row in range(a.shape[0]):
        for column in range(b.shape[1]):
            for shared in range(a.shape[1]):
                result[row, column] += a[row, shared] * b[shared, column]
    return result


def stable_softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    shifted = x - np.max(x, axis=axis, keepdims=True)
    exponentials = np.exp(shifted)
    return exponentials / exponentials.sum(axis=axis, keepdims=True)


def main() -> None:
    a = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
    b = np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
    expected = np.array([[22.0, 28.0], [49.0, 64.0]])

    product = manual_matmul(a, b)
    assert np.allclose(product, expected)
    assert np.allclose(product, a @ b)

    # Matrix multiplication defines a linear map: addition and scalar
    # multiplication can happen before or after the map with the same result.
    matrix = np.array([[2.0, -1.0], [0.5, 3.0]])
    x = np.array([1.0, 2.0])
    y = np.array([-2.0, 4.0])
    scalar = 2.5
    assert np.allclose((x + y) @ matrix, x @ matrix + y @ matrix)
    assert np.allclose((scalar * x) @ matrix, scalar * (x @ matrix))

    # Adding a nonzero bias preserves neither f(0)=0 nor strict linearity.
    affine_bias = np.array([1.0, -2.0])
    assert not np.allclose(np.zeros(2) @ matrix + affine_bias, np.zeros(2))

    batch = np.arange(12.0).reshape(4, 3)
    bias = np.array([10.0, 20.0, 30.0])
    broadcast_result = batch + bias
    assert broadcast_result.shape == (4, 3)

    dependent = np.array([[1.0, 2.0], [2.0, 4.0]])
    assert np.linalg.matrix_rank(dependent) == 1

    system = np.array([[3.0, 1.0], [1.0, 2.0]])
    target = np.array([9.0, 8.0])
    solution = np.linalg.solve(system, target)
    assert np.allclose(solution, [2.0, 3.0])
    assert np.allclose(system @ solution, target)

    probabilities = stable_softmax(np.array([[1000.0, 1001.0, 1002.0]]))
    assert np.allclose(probabilities.sum(axis=-1), 1.0)

    print("A @ B =\n", product)
    print("linearity checks: addition and scaling preserved")
    print("broadcast shape:", broadcast_result.shape)
    print("rank:", np.linalg.matrix_rank(dependent))
    print("solution to Ax=b:", solution)
    print("stable probabilities:", probabilities)

    # Try it yourself:
    # 1. Vectorize manual_matmul without calling np.matmul or @.
    # 2. Add batched inputs with shape [B, M, N] and [B, N, P].
    # 3. Demonstrate a silent but incorrect broadcasting operation.


if __name__ == "__main__":
    main()