"""Small typed experiments for the foundations course. Run this file directly."""
from __future__ import annotations

import numpy as np
from numpy.typing import NDArray

FloatArray = NDArray[np.float64]

# region shapes
def project(x: FloatArray, weight: FloatArray, bias: FloatArray) -> FloatArray:
    """x: (tokens, input_features), weight: (input_features, output_features)."""
    if x.ndim != 2 or weight.ndim != 2 or x.shape[1] != weight.shape[0]:
        raise ValueError("The input feature axis must match the weight's row axis")
    products = x @ weight
    output = products + bias  # (output_features,) broadcasts over tokens.
    return output
# endregion shapes

# region mixture
def weighted_mix(weights: FloatArray, values: FloatArray) -> FloatArray:
    """weights: (sources,), values: (sources, features); result: (features,)."""
    if np.any(weights < 0) or not np.isclose(weights.sum(), 1.0):
        raise ValueError("Use nonnegative weights summing to one")
    contributions = weights[:, None] * values
    output = contributions.sum(axis=0)
    return output
# endregion mixture

# region probability
def softmax(scores: FloatArray, temperature: float = 1.0) -> FloatArray:
    """Normalize the final axis; temperature must be positive."""
    if temperature <= 0:
        raise ValueError("Temperature must be positive")
    scaled = scores / temperature
    shifted = scaled - scaled.max(axis=-1, keepdims=True)
    evidence = np.exp(shifted)
    probabilities = evidence / evidence.sum(axis=-1, keepdims=True)
    return probabilities
# endregion probability

# region learning
def learn_step(weight: float, x: float, target: float, rate: float = 0.1) -> tuple[float, float, float]:
    """One squared-error gradient step for prediction = weight * x."""
    prediction = weight * x
    error = prediction - target
    loss = error ** 2
    gradient = 2.0 * error * x
    updated = weight - rate * gradient
    return updated, loss, gradient
# endregion learning


def check() -> None:
    x = np.array([[1.0, 2.0], [3.0, 4.0]])
    weight = np.array([[1.0, 0.0], [0.5, 2.0]])
    np.testing.assert_allclose(project(x, weight, np.array([0.0, 1.0])), [[2, 5], [5, 9]])
    np.testing.assert_allclose(weighted_mix(np.array([0.2, 0.3, 0.5]), np.array([[1., 0.], [0., 2.], [2., 1.]])), [1.2, 1.1])
    np.testing.assert_allclose(softmax(np.array([0., np.log(2.), np.log(3.)])), [1/6, 2/6, 3/6])
    np.testing.assert_allclose(learn_step(0.5, 2.0, 3.0), [1.3, 4.0, -8.0])
    print("PASS: foundations projections, mixtures, softmax, and gradient update")


if __name__ == "__main__":
    check()