import copy
import torch
from torch import nn

# The same experiment design; PyTorch and browser RNGs differ.
g = torch.Generator().manual_seed(17)
x = 2 * torch.rand(100, 1, generator=g) - 1
y = torch.sin(3 * x) + 0.3 * x + (2 * torch.rand(x.shape, generator=g) - 1) * 0.3
x_train, y_train = x[:20], y[:20]
x_val, y_val = x[20:], y[20:]

def run(dropout=0.0, decay=0.0, early_stopping=False):
    torch.manual_seed(7)  # Paired initial weights.
    model = nn.Sequential(
        nn.Linear(1, 24), nn.Tanh(), nn.Dropout(dropout),
        nn.Linear(24, 24), nn.Tanh(), nn.Dropout(dropout),
        nn.Linear(24, 1),
    )
    for layer in model.modules():
        if isinstance(layer, nn.Linear):
            nn.init.xavier_uniform_(layer.weight)
            nn.init.zeros_(layer.bias)
    weights = [p for name, p in model.named_parameters() if name.endswith("weight")]
    biases = [p for name, p in model.named_parameters() if name.endswith("bias")]
    optimizer = torch.optim.AdamW([
        {"params": weights, "weight_decay": decay},
        {"params": biases, "weight_decay": 0.0},
    ], lr=0.01)
    best, wait, best_epoch = float("inf"), 0, 0
    best_state = copy.deepcopy(model.state_dict())
    for epoch in range(1, 601):
        model.train()  # Dropout active only here.
        optimizer.zero_grad()
        loss = nn.functional.mse_loss(model(x_train), y_train)
        loss.backward()
        optimizer.step()
        model.eval()
        with torch.inference_mode():
            val_rmse = ((model(x_val) - y_val) ** 2).mean().sqrt().item()
        if val_rmse < best - 1e-5:
            best, wait, best_epoch = val_rmse, 0, epoch
            best_state = copy.deepcopy(model.state_dict())
        else:
            wait += 1
        if early_stopping and wait >= 60:
            break
    if early_stopping:
        model.load_state_dict(best_state)
    model.eval()
    with torch.inference_mode():
        train_rmse = ((model(x_train) - y_train) ** 2).mean().sqrt().item()
        val_rmse = ((model(x_val) - y_val) ** 2).mean().sqrt().item()
    return {"train_RMSE": train_rmse, "validation_RMSE": val_rmse,
            "epochs_run": epoch, "selected_epoch": best_epoch if early_stopping else epoch}

print("Control:", run())
print("Experiment:", run(dropout=0, decay=0,
                         early_stopping=True))
# Validation guides choices. Keep a separate final test set untouched.