"""A small network: fit a curve that no straight line can follow.

Run: python3 examples/pytorch_small_network.py
Uses the training loop from pytorch_first_update.py unchanged; only the model (the network) is new.
"""

import torch
from torch import nn

x = torch.linspace(-1, 1, 41).reshape(-1, 1)   # 41 inputs, one per row: shape (41, 1)
y = 1 - (x - 0.2) ** 2                         # the curve we want the network to learn


def make_network(hidden):
    if hidden == 0:
        return nn.Linear(1, 1)                 # no hidden layer: the straight line from PyTorch basics
    return nn.Sequential(nn.Linear(1, hidden), nn.ReLU(), nn.Linear(hidden, 1))


def train(hidden, seed, steps=3000):
    torch.manual_seed(seed)                    # the random starting weights
    network = make_network(hidden)
    loss_fn = nn.MSELoss()
    optimizer = torch.optim.SGD(network.parameters(), lr=0.1)
    for _ in range(steps):
        loss = loss_fn(network(x), y)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    return network, loss_fn(network(x), y).item()


if __name__ == "__main__":
    layer = nn.Linear(4, 16)
    print("nn.Linear(4, 16) weight", tuple(layer.weight.shape), "bias", tuple(layer.bias.shape))
    print("ReLU of [-2, -0.5, 0, 1.5]:", torch.relu(torch.tensor([-2.0, -0.5, 0.0, 1.5])).tolist())

    network = nn.Sequential(nn.Linear(1, 4), nn.ReLU(), nn.Linear(4, 1))
    first, relu, second = network               # nn.Sequential holds the three pieces in order
    print("network(x) == second(relu(first(x))):", torch.equal(network(x), second(relu(first(x)))))
    print("first.weight", tuple(first.weight.shape), "first.bias", tuple(first.bias.shape))   # one row per hidden neuron

    for hidden in (0, 1, 2, 4, 16):
        losses = [train(hidden, seed)[1] for seed in (0, 1, 2)]
        print(f"{hidden:2d} hidden neurons, starts 1-3: loss " + "  ".join(f"{v:.4f}" for v in losses))