import torch
import torch.nn as nn
import torch.nn.functional as F


class RegularizedMLP(nn.Module):
    def __init__(self, d_in, d_hidden, d_out, dropout_p=0.5):
        super().__init__()
        self.fc1 = nn.Linear(d_in, d_hidden)
        self.bn1 = nn.BatchNorm1d(d_hidden)
        self.dropout1 = nn.Dropout(dropout_p)
        self.fc2 = nn.Linear(d_hidden, d_hidden)
        self.bn2 = nn.BatchNorm1d(d_hidden)
        self.dropout2 = nn.Dropout(dropout_p)
        self.fc3 = nn.Linear(d_hidden, d_out)

    def forward(self, x):
        x = F.relu(self.bn1(self.fc1(x)))
        x = self.dropout1(x)
        x = F.relu(self.bn2(self.fc2(x)))
        x = self.dropout2(x)
        return self.fc3(x)


def train_with_early_stopping(
    model, train_loader, val_loader, epochs=100,
    lr=1e-3, weight_decay=1e-4, patience=10,
):
    optimizer = torch.optim.Adam(
        model.parameters(), lr=lr, weight_decay=weight_decay
    )
    best_val_loss = float("inf")
    best_state = None
    epochs_without_improvement = 0

    for epoch in range(epochs):
        model.train()
        for xb, yb in train_loader:
            optimizer.zero_grad()
            loss = F.cross_entropy(model(xb), yb)
            loss.backward()
            optimizer.step()

        model.eval()
        val_loss = 0.0
        with torch.no_grad():
            for xb, yb in val_loader:
                val_loss += F.cross_entropy(
                    model(xb), yb, reduction="sum"
                ).item()
        val_loss /= len(val_loader.dataset)

        if val_loss < best_val_loss:
            best_val_loss = val_loss
            best_state = {
                k: v.clone() for k, v in model.state_dict().items()
            }
            epochs_without_improvement = 0
        else:
            epochs_without_improvement += 1
            if epochs_without_improvement >= patience:
                print(f"Early stopping at epoch {epoch}")
                break

    model.load_state_dict(best_state)
    return model, best_val_loss
