"""Tensor versions of the foundations experiments, with automatic gradients."""
from __future__ import annotations

import torch
from torch import Tensor

# region shapes
def project(x: Tensor, weight: Tensor, bias: Tensor) -> Tensor:
    """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
    return output
# endregion shapes

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

# region probability
def softmax(scores: Tensor, temperature: float = 1.0) -> Tensor:
    if temperature <= 0:
        raise ValueError("Temperature must be positive")
    scaled = scores / temperature
    probabilities = scaled.softmax(dim=-1)
    return probabilities
# endregion probability

# region learning
def learn_step(weight: float, x: float, target: float, rate: float = 0.1) -> tuple[float, float, float]:
    parameter = torch.tensor(weight, dtype=torch.float64, requires_grad=True)
    prediction = parameter * x
    loss = (prediction - target) ** 2
    loss.backward()
    assert parameter.grad is not None
    gradient = parameter.grad.item()
    updated = parameter.item() - rate * gradient
    return updated, loss.item(), gradient
# endregion learning


def check() -> None:
    import foundations_numpy as reference
    import numpy as np

    x = np.array([[1., 2.], [3., 4.]])
    w = np.array([[1., 0.], [0.5, 2.]])
    b = np.array([0., 1.])
    np.testing.assert_allclose(project(torch.tensor(x), torch.tensor(w), torch.tensor(b)).numpy(), reference.project(x, w, b))
    np.testing.assert_allclose(learn_step(0.5, 2.0, 3.0), reference.learn_step(0.5, 2.0, 3.0))
    print("PASS: foundations NumPy/PyTorch agreement")


if __name__ == "__main__":
    check()