"""A complete PyTorch update with one input and one adjustable model.

Run: python3 examples/pytorch_first_update.py
The matching browser lesson starts from these exact numbers.
"""

import torch
from torch import nn


# The model is guess = weight × input + bias.
model = nn.Linear(1, 1)
with torch.no_grad():
    model.weight.fill_(1.0)
    model.bias.fill_(0.0)

input_value = torch.tensor([[2.0]])
answer = torch.tensor([[5.0]])
loss_fn = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.05)

guess_before = model(input_value)
loss_before = loss_fn(guess_before, answer)
optimizer.zero_grad()  # Remove notes from an earlier example, if any.
loss_before.backward()  # Write a gradient beside each parameter; do not move it yet.

print(f"Guess before: {guess_before.item():.2f}; answer: {answer.item():.2f}")
print(f"Squared error: {loss_before.item():.2f}")
print(f"Gradient for weight: {model.weight.grad.item():.2f}")
print(f"Gradient for bias: {model.bias.grad.item():.2f}")

optimizer.step()  # This is the line that moves the parameters.
with torch.no_grad():
    guess_after = model(input_value)
    loss_after = loss_fn(guess_after, answer)

print(f"Weight: {model.weight.item():.2f}; bias: {model.bias.item():.2f}")
print(f"Guess after: {guess_after.item():.2f}; squared error: {loss_after.item():.2f}")