# Source: content/notes/deep-learning/neural-networks.md
# Independent CPU example; use the curriculum environment.
# See /notes/ml/#example-environment or /notes/deep-learning/#example-environment.

import copy
import numpy as np
import torch
from torch import nn
from sklearn.datasets import make_moons
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

torch.set_num_threads(1)
torch.manual_seed(17)
np.random.seed(17)
X, y = make_moons(n_samples=1000, noise=0.16, random_state=17)
X_train, X_hold, y_train, y_hold = train_test_split(
    X, y, test_size=0.4, stratify=y, random_state=18
)
X_val, X_test, y_val, y_test = train_test_split(
    X_hold, y_hold, test_size=0.5, stratify=y_hold, random_state=19
)
scaler = StandardScaler().fit(X_train)
def tensors(features, labels):
    return (torch.tensor(scaler.transform(features), dtype=torch.float32),
            torch.tensor(labels, dtype=torch.long))
xt, yt = tensors(X_train, y_train)
xv, yv = tensors(X_val, y_val)
xs, ys = tensors(X_test, y_test)
loss_fn = nn.CrossEntropyLoss()

def train(hidden):
    torch.manual_seed(20)
    model = (nn.Sequential(nn.Linear(2, 24), nn.Tanh(), nn.Linear(24, 2))
             if hidden else nn.Linear(2, 2))
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.02, weight_decay=0.001)
    best_loss = float("inf")
    best_state = copy.deepcopy(model.state_dict())
    for epoch in range(300):
        model.train()
        optimizer.zero_grad(set_to_none=True)
        loss = loss_fn(model(xt), yt)
        assert torch.isfinite(loss)
        loss.backward()
        optimizer.step()
        model.eval()
        with torch.inference_mode():
            val_loss = loss_fn(model(xv), yv).item()
        if val_loss < best_loss:
            best_loss = val_loss
            best_state = copy.deepcopy(model.state_dict())
    model.load_state_dict(best_state)
    model.eval()
    with torch.inference_mode():
        test_loss = loss_fn(model(xs), ys).item()
        accuracy = (model(xs).argmax(1) == ys).float().mean().item()
    print("MLP" if hidden else "affine", "validation", round(best_loss, 4),
          "test loss", round(test_loss, 4), "test accuracy", round(accuracy, 3))
    return model, accuracy

linear, linear_accuracy = train(False)
mlp, mlp_accuracy = train(True)
assert mlp_accuracy > 0.92
assert mlp_accuracy > linear_accuracy + 0.04

# Verify a smooth two-layer backward pass in float64.
rng = np.random.default_rng(21)
a = rng.normal(size=(5, 3))
w1 = rng.normal(size=(3, 4)) * 0.2
b1 = rng.normal(size=4) * 0.1
w2 = rng.normal(size=(4, 2)) * 0.2
b2 = np.zeros(2)
labels = np.array([0, 1, 0, 1, 1])
h = np.tanh(a @ w1 + b1)
logits = h @ w2 + b2
shifted = logits - logits.max(axis=1, keepdims=True)
log_probs = shifted - np.log(np.exp(shifted).sum(axis=1, keepdims=True))
loss_np = -log_probs[np.arange(5), labels].mean()
g = np.exp(log_probs)
g[np.arange(5), labels] -= 1
g /= len(labels)
gw2, gb2 = h.T @ g, g.sum(axis=0)
g1 = (g @ w2.T) * (1 - h * h)
gw1, gb1 = a.T @ g1, g1.sum(axis=0)
params = [torch.tensor(v, dtype=torch.float64, requires_grad=True)
          for v in (w1, b1, w2, b2)]
tw1, tb1, tw2, tb2 = params
z = torch.tanh(torch.tensor(a) @ tw1 + tb1) @ tw2 + tb2
loss_torch = nn.functional.cross_entropy(z, torch.tensor(labels))
loss_torch.backward()
np.testing.assert_allclose(loss_np, loss_torch.item(), rtol=1e-12)
for parameter, expected in zip(params, (gw1, gb1, gw2, gb2)):
    np.testing.assert_allclose(parameter.grad.numpy(), expected, atol=1e-12)
print("All manual gradients match autograd.")