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

import copy
import math
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(62)
np.random.seed(62)
X, y = make_moons(n_samples=600, noise=0.25, random_state=62)
train_x, hold_x, train_y, hold_y = train_test_split(
    X, y, train_size=100, stratify=y, random_state=63)
val_x, test_x, val_y, test_y = train_test_split(
    hold_x, hold_y, train_size=100, stratify=hold_y, random_state=64)
scaler = StandardScaler().fit(train_x)
def convert(features, labels):
    return (torch.tensor(scaler.transform(features), dtype=torch.float32),
            torch.tensor(labels, dtype=torch.long))
xt, yt = convert(train_x, train_y)
xv, yv = convert(val_x, val_y)
xs, ys = convert(test_x, test_y)
trials = []
for dropout_rate, decay in ((0.0, 0.0), (0.2, 0.01), (0.4, 0.1)):
    torch.manual_seed(65)
    model = nn.Sequential(nn.Linear(2, 32), nn.ReLU(), nn.Dropout(dropout_rate),
                          nn.Linear(32, 32), nn.ReLU(), nn.Linear(32, 2))
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.01, weight_decay=decay)
    best, best_epoch = float("inf"), -1
    state = copy.deepcopy(model.state_dict())
    for epoch in range(180):
        model.train()
        optimizer.zero_grad(set_to_none=True)
        loss = nn.functional.cross_entropy(model(xt), yt)
        assert torch.isfinite(loss)
        loss.backward()
        optimizer.step()
        model.eval()
        with torch.inference_mode():
            validation = nn.functional.cross_entropy(model(xv), yv).item()
        if validation < best:
            best, best_epoch = validation, epoch
            state = copy.deepcopy(model.state_dict())
    model.load_state_dict(state)
    model.eval()
    with torch.inference_mode():
        training = nn.functional.cross_entropy(model(xt), yt).item()
    trials.append((best, model, dropout_rate, decay))
    print("dropout/decay", dropout_rate, decay, "best epoch", best_epoch,
          "training/validation", training, best)
selected = min(trials, key=lambda item: item[0])
best, model, dropout_rate, decay = selected
with torch.inference_mode():
    test_loss = nn.functional.cross_entropy(model(xs), ys).item()
    test_accuracy = (model(xs).argmax(1) == ys).float().mean().item()
assert best < math.log(2)
assert test_accuracy > 0.8
print("Selected:", dropout_rate, decay, "test loss/accuracy:", test_loss, test_accuracy)