"""Lesson 1 lab: linear transformations and tiny semantic search.

Run:
    python 01-linear-algebra-lab.py

Requires:
    numpy
"""

from __future__ import annotations

import numpy as np


def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """Return cosine similarity and reject zero vectors explicitly."""
    denominator = np.linalg.norm(a) * np.linalg.norm(b)
    if denominator == 0:
        raise ValueError("Cosine similarity is undefined for a zero vector.")
    return float(a @ b / denominator)


def transform(vector: np.ndarray, matrix: np.ndarray) -> np.ndarray:
    """Apply a row-vector transformation after validating the shape contract."""
    if vector.ndim != 1 or matrix.ndim != 2:
        raise ValueError("Expected a 1D vector and a 2D matrix.")
    if vector.shape[0] != matrix.shape[0]:
        raise ValueError(f"Cannot multiply {vector.shape} by {matrix.shape}.")
    return vector @ matrix


def affine_transform(
    vector: np.ndarray, matrix: np.ndarray, bias: np.ndarray
) -> np.ndarray:
    """Apply xW+b and make the affine offset explicit."""
    transformed = transform(vector, matrix)
    if bias.shape != transformed.shape:
        raise ValueError(f"Bias {bias.shape} must match output {transformed.shape}.")
    return transformed + bias


def rank_documents(
    query: np.ndarray, documents: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
    """Rank document rows by cosine similarity to one query."""
    if documents.ndim != 2 or query.shape != (documents.shape[1],):
        raise ValueError("Query dimension must match document columns.")

    document_norms = np.linalg.norm(documents, axis=1, keepdims=True)
    if np.any(document_norms == 0) or np.linalg.norm(query) == 0:
        raise ValueError("Zero embeddings require an explicit product policy.")

    normalized_documents = documents / document_norms
    normalized_query = query / np.linalg.norm(query)
    scores = normalized_documents @ normalized_query
    order = np.argsort(scores)[::-1]
    return order, scores[order]


def main() -> None:
    vector = np.array([2.0, 3.0])
    scaling = np.array([[2.0, 0.0], [0.0, 0.5]])
    transformed = transform(vector, scaling)
    assert np.allclose(transformed, [4.0, 1.5])

    # A matrix map is linear: it maps zero to zero. A nonzero bias makes the
    # familiar neural-network "linear layer" affine rather than strictly linear.
    zero = np.zeros(2)
    bias = np.array([1.0, -1.0])
    assert np.allclose(transform(zero, scaling), zero)
    assert np.allclose(affine_transform(zero, scaling, bias), bias)

    labels = np.array(["neural networks", "database indexes", "gradient descent"])
    documents = np.array(
        [[0.95, 0.80, 0.10], [0.10, 0.20, 0.95], [0.90, 0.85, 0.05]]
    )
    query = np.array([0.92, 0.82, 0.08])
    order, ranked_scores = rank_documents(query, documents)

    print("x =", vector, "xW =", transformed)
    print("affine zero image (0W+b) =", affine_transform(zero, scaling, bias))
    print("cos([1, 1], [2, 2]) =", cosine_similarity(np.array([1, 1]), np.array([2, 2])))
    print("\nSemantic ranking")
    for index, score in zip(order, ranked_scores, strict=True):
        print(f"{score:.3f}  {labels[index]}")

    # Try it yourself:
    # 1. Add a second query and rank both with one matrix multiplication.
    # 2. Create a rotation matrix and inspect how it changes a vector.
    # 3. Trigger each validation error and explain why it exists.


if __name__ == "__main__":
    main()