"""Train y=2x+1 and check autograd, shapes, and soft-label cross entropy. Requires torch."""
from __future__ import annotations
import torch
from torch import nn
from torch.nn import functional as F


def fit_line(*, epochs=300):
    """Pre: epochs>0. Post: trained model and loss trace on deterministic synthetic data."""
    if epochs <= 0:
        raise ValueError("epochs must be positive")
    torch.manual_seed(7)
    x = torch.linspace(-1, 1, 64).reshape(-1, 1)
    y = 2*x + 1
    model = nn.Linear(1, 1)
    optimizer = torch.optim.SGD(model.parameters(), lr=.05)
    history = []
    # Invariant: completed steps have finite recorded losses. Bound: epochs-step.
    for _ in range(epochs):
        optimizer.zero_grad(set_to_none=True)
        predictions = model(x)
        loss = F.mse_loss(predictions, y)
        loss.backward()
        optimizer.step()
        history.append(loss.item())
    return model, history


def demo():
    torch.set_num_threads(1)
    model, history = fit_line()
    model.eval()
    with torch.no_grad():
        prediction = model(torch.tensor([[.5]])).item()
    print("Initial/final MSE:", history[0], history[-1])
    print("Learned w,b:", model.weight.item(), model.bias.item())
    print("Prediction for x=.5:", prediction, "; target=2")
    logits = torch.tensor([[2., 1., 0.]])
    soft_target = torch.tensor([[.8, .1, .1]])
    print("CrossEntropy accepts float soft targets:", F.cross_entropy(logits, soft_target).item())
    print("A shape is a contract: inputs and regression targets here are both (64,1).")


def self_test():
    torch.set_num_threads(1)
    x = torch.tensor([2.], requires_grad=True)
    (x*x).sum().backward()
    check(x.grad.item(), 4.)
    model, history = fit_line()
    check(history[-1] < 1e-7, True)
    check(history[-1] < history[0], True)
    check(model(torch.zeros(3, 1)).shape, torch.Size([3, 1]))
    z = torch.tensor([[2., 1., 0.]])
    target = torch.tensor([[.8, .1, .1]])
    expected = -(target*F.log_softmax(z, dim=-1)).sum()
    check(torch.allclose(F.cross_entropy(z, target), expected), True)
    return 5

def check(actual, expected):
    """Require an expected result; this check also runs under python -O."""
    if actual != expected:
        raise AssertionError(f"expected {expected!r}, got {actual!r}")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        print(f"{self_test()} checks passed")
    else:
        demo()