"""Which way to correct each parameter: nudges, the chain rule, PyTorch's backward(), and a batch.

Run: python3 examples/pytorch_gradients.py
The model is the one from PyTorch basics: guess = weight * input + bias, loss = squared miss.
"""

import torch
from torch import nn


def loss_at(w, b, x, answer):
    return (w * x + b - answer) ** 2


# 1. Nudge each parameter a little and watch the loss.
x, answer, w, b = 2.0, 5.0, 1.0, 0.0
base = loss_at(w, b, x, answer)
print(f"guess {w * x + b:.2f}, loss {base:.4f}")
for name, nudged in [("w", loss_at(w + 0.01, b, x, answer)), ("b", loss_at(w, b + 0.01, x, answer))]:
    print(f"  nudge {name} up by 0.01: loss {nudged:.4f}, change {nudged - base:+.4f}, rate {(nudged - base) / 0.01:+.2f}")

# 2. The chain rule by hand: how the loss reacts to the guess, times how the guess reacts to each parameter.
guess = w * x + b
dloss_dguess = 2 * (guess - answer)
print(f"chain: dloss/dguess = {dloss_dguess:+.0f}; dloss/dw = {dloss_dguess:+.0f} x {x:.0f} = {dloss_dguess * x:+.0f}; "
      f"dloss/db = {dloss_dguess:+.0f} x 1 = {dloss_dguess:+.0f}")

# 3. PyTorch's backward() finds the same numbers, for any input.
loss_fn = nn.MSELoss()
for x in (2.0, -2.0):
    model = nn.Linear(1, 1)
    with torch.no_grad():
        model.weight.fill_(1.0)
        model.bias.fill_(0.0)
    loss = loss_fn(model(torch.tensor([[x]])), torch.tensor([[5.0]]))
    loss.backward()
    print(f"backward() with input {x:+.0f}: weight gradient {model.weight.grad.item():+.0f}, "
          f"bias gradient {model.bias.grad.item():+.0f}")

# 4. A batch: each example's own gradient (its vote), and backward() on the batch's mean loss.
inputs = torch.tensor([[1.0], [2.0], [3.0]])
answers = torch.tensor([[3.0], [5.0], [7.0]])
model = nn.Linear(1, 1)
with torch.no_grad():
    model.weight.fill_(1.0)
    model.bias.fill_(0.0)
for i in range(3):
    model.zero_grad()
    loss_fn(model(inputs[i:i + 1]), answers[i:i + 1]).backward()
    print(f"example {i + 1}: vote for weight {model.weight.grad.item():+.2f}, for bias {model.bias.grad.item():+.2f}")
model.zero_grad()
loss_fn(model(inputs), answers).backward()      # the mean loss over all three rows
print(f"whole batch:  weight {model.weight.grad.item():+.2f}, bias {model.bias.grad.item():+.2f}  (the average of the votes)")

# 5. Backpropagation in a tiny network: x -> z = w1*x + b1 -> h = ReLU(z) -> y = v*h + c -> loss.
def tiny(w1, b1, v, c):
    net = nn.Sequential(nn.Linear(1, 1), nn.ReLU(), nn.Linear(1, 1))
    with torch.no_grad():
        net[0].weight.fill_(w1); net[0].bias.fill_(b1)
        net[2].weight.fill_(v); net[2].bias.fill_(c)
    return net


def grads(net):
    return [p.grad.item() for p in (net[0].weight, net[0].bias, net[2].weight, net[2].bias)]


x, answer = torch.tensor([[2.0]]), torch.tensor([[5.0]])
for label, start in [("neuron on ", (1.0, 0.0, 1.0, 0.0)), ("neuron off", (1.0, -3.0, 1.0, 0.0))]:
    net = tiny(*start)
    loss = loss_fn(net(x), answer)                  # forward: PyTorch records every operation
    loss.backward()                                 # one sweep: a gradient for every parameter
    print(f"{label}: loss {loss.item():.2f}, gradients w1 b1 v c = " + " ".join(f"{g:+.0f}" for g in grads(net)))

# Check the sweep against nudging each parameter (neuron on), then update all four at once.
net = tiny(1.0, 0.0, 1.0, 0.0)
base = loss_fn(net(x), answer).item()
nudged = []
for p in net.parameters():
    with torch.no_grad():
        p += 0.01
        nudged.append((loss_fn(net(x), answer).item() - base) / 0.01)
        p -= 0.01
print("nudge estimates w1 b1 v c = " + " ".join(f"{g:+.2f}" for g in nudged))

optimizer = torch.optim.SGD(net.parameters(), lr=0.01)
losses = []
for iteration in range(20):
    loss = loss_fn(net(x), answer)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    losses.append(loss.item())
    if iteration == 0:
        print("after one update: w1 b1 v c = " + " ".join(f"{p.item():.2f}" for p in net.parameters()),
              f"-> new loss {loss_fn(net(x), answer).item():.2f}")
print("loss over 20 iterations:", " ".join(f"{v:.3f}" for v in losses))