import math
import torch
import random
import torch.nn as nn
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader, TensorDataset



#-------------------------------------------------- SEED EVERYTHING --------------------------------------------------
def seed_everything(seed):
    '''Set the seed for all random generators to be the provided seed'''
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)



#-------------------------------------------------- CREATE DATASETS --------------------------------------------------
#------------------------- CLASSIFICATION DATASET -------------------------
def create_classification_dataset(seed=0):
    seed_everything(seed)

    d = 10
    I = 2*math.pi
    n_train, n_test = 2000, 500

    w1, w2 = torch.randn(d)*I, torch.randn(d)*I
    # Fixed random vectors w1, w2 ~ N(0, I)

    X_train = torch.randn(n_train,d)
    X_test = torch.randn(n_test,d)
    # Each feature independently sampled from N(0,1)

    y_train = (torch.sin(X_train @ w1) + 0.5 * (X_train @ w2) + torch.randn(n_train)*0.1 > 0).long()
    y_test = (torch.sin(X_test @ w1) + 0.5 * (X_test @ w2) + torch.randn(n_test)*0.1 > 0).long()
    # y = I( sin(w1^T x) + 0.5 w2^T x + ε > 0 )

    mean = X_train.mean(dim=0, keepdim=True)
    std = X_train.std(dim=0, keepdim=True)

    X_train = (X_train - mean) / std
    X_test = (X_test - mean) / std
    # Training statistics only, as per required.

    return X_train, X_test, y_train, y_test



#------------------------- REGRESSION DATASET -------------------------
def create_regression_dataset(seed=0):
    seed_everything(seed)

    d = 10
    I = 2*math.pi
    n_train, n_test = 2000, 500

    w = torch.randn(d)*I

    X_train = torch.randn(n_train,d)
    X_test = torch.randn(n_test,d)

    y_train = torch.sin(X_train @ w) + 0.1 * (torch.norm(X_train, dim=1)**2) + torch.randn(n_train)*0.1
    y_test = torch.sin(X_test @ w) + 0.1 * (torch.norm(X_test, dim=1)**2) + torch.randn(n_test)*0.1
    # y = sin(w^T x) + 0.1||x||^2 + ε

    mean = X_train.mean(dim=0, keepdim=True)
    std = X_train.std(dim=0, keepdim=True)

    X_train = (X_train - mean) / std
    X_test = (X_test - mean) / std

    return X_train, X_test, y_train, y_test



#-------------------------------------------------- MLP IMPLEMENTATION --------------------------------------------------
class MLP(nn.Module):
    def __init__(self, input_dim, hidden_layers, output_dim, activation="relu"):
        super(MLP, self).__init__()

        match activation:
            case "relu": self.activation = nn.ReLU()
            case "sigmoid": self.activation = nn.Sigmoid()
            case "tanh": self.activation = nn.Tanh()

        layers = []
        prev_dim = input_dim

        #------------------------- LAYER NETWORK -------------------------
        for hidden_dim in hidden_layers:
            layers.append(nn.Linear(prev_dim, hidden_dim))
            layers.append(self.activation)
            prev_dim = hidden_dim
        
        layers.append(nn.Linear(prev_dim, output_dim))

        self.model = nn.Sequential(*layers)

        self._initialize_weights(activation)



    #------------------------- WEIGHT INITIALIZATION -------------------------
    def _initialize_weights(self, activation):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                match activation:
                    case "relu": nn.init.kaiming_normal_(m.weight, nonlinearity='relu')
                    case "sigmoid" | "tanh": nn.init.xavier_normal_(m.weight)
                nn.init.zeros_(m.bias)
        # He for ReLU, Xavier for Sigmoid/Tanh



    #------------------------- FORWARD -------------------------
    def forward(self, x):
        return self.model(x)



    #------------------------- PREDICTION FUNCTION -------------------------
    def predict(self, x):
        self.eval()
        with torch.no_grad():
            output = self.forward(x)
            if output.shape[-1] == 1:
                return (torch.sigmoid(output) > 0.5).long()
            return torch.argmax(output, dim=1)



    #------------------------- PARAMETER COUNT -------------------------
    def count_parameters(self):
        return sum(p.numel() for p in self.parameters() if p.requires_grad)



#-------------------------------------------------- TRAINING LOOP --------------------------------------------------
def train(model, X_train, y_train, X_val, y_val, task="classification", optimizer_name="adam", lr=1e-3, batch_size=64, epochs=50, weight_decay=0.0, momentum=0.9, tol=None):
    train_dataset = TensorDataset(X_train, y_train)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    val_dataset = TensorDataset(X_val, y_val)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

    match task:
        case "classification": criterion = nn.BCEWithLogitsLoss()
        case "regression": criterion = nn.MSELoss()

    match optimizer_name.upper():
        case "SGD": opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum, weight_decay=weight_decay)
        case "ADAM": opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)

    train_losses, val_losses = [], []
    convergence_epoch = None

    #------------------------- EPOCH LOOP -------------------------
    for epoch in range(epochs):
        model.train()
        total_train_loss = 0

        for xb, yb in train_loader:
            if task == "classification":
                yb = yb.float().unsqueeze(1)

            outputs = model(xb)
            loss = criterion(outputs, yb)

            opt.zero_grad()
            loss.backward()
            opt.step()

            total_train_loss += loss.item()

        avg_train_loss = total_train_loss / len(train_loader)
        train_losses.append(avg_train_loss)

        #------------------------- VALIDATION LOOP -------------------------
        model.eval()
        total_val_loss = 0

        with torch.no_grad():
            for xb, yb in val_loader:
                if task == "classification":
                    yb = yb.float().unsqueeze(1)

                outputs = model(xb)
                loss = criterion(outputs, yb)
                total_val_loss += loss.item()

        avg_val_loss = total_val_loss / len(val_loader)
        val_losses.append(avg_val_loss)

        if tol is not None and convergence_epoch is None and avg_train_loss < tol:
            convergence_epoch = epoch + 1

    return train_losses, val_losses, convergence_epoch



#-------------------------------------------------- EXPERIMENT A: ACTIVATIONS --------------------------------------------------
def experiment_activation_functions():
    X_train, X_val, y_train, y_val = create_classification_dataset()

    activations = ["relu", "sigmoid", "tanh"]
    loss_dict = {}

    for act in activations:
        model = MLP(10, [64, 32], 1, activation=act)

        train_losses, _, _ = train(model, X_train, y_train, X_val, y_val)

        loss_dict[act.upper()] = train_losses

    plot(loss_dict, "Experiment A: Activation Function Comparison")



#-------------------------------------------------- EXPERIMENT B: DEPTH VS WIDTH --------------------------------------------------
def experiment_depth_vs_width():
    X_train, X_val, y_train, y_val = create_classification_dataset()

    architectures = {"Shallow-Wide (50)": [50], "Deep-Narrow (20,20,20)": [20,20,20]}

    loss_dict = {}

    for name, hidden in architectures.items():
        model = MLP(10, hidden, 1, activation="relu")

        train_losses, _, _ = train(model, X_train, y_train, X_val, y_val)

        loss_dict[name] = train_losses

    plot(loss_dict, "Experiment B: Depth vs Width")



#-------------------------------------------------- EXPERIMENT C: OPTIMIZERS --------------------------------------------------
def experiment_optimizers():
    X_train, X_val, y_train, y_val = create_classification_dataset()

    optimizers = ["SGD", "ADAM"]
    loss_dict = {}

    for opt in optimizers:
        model = MLP(10, [64,32], 1, activation="relu")

        train_losses, _, _ = train(model, X_train, y_train, X_val, y_val, optimizer_name=opt)

        loss_dict[opt] = train_losses

    plot(loss_dict, "Experiment C: Optimizer Comparison")



#-------------------------------------------------- EXPERIMENT D: REGULARIZATION --------------------------------------------------
def experiment_regularization():
    X_train, X_val, y_train, y_val = create_classification_dataset()

    settings = {"No L2": 0.0, "With L2 (1e-3)": 1e-3}

    loss_dict = {}

    for name, wd in settings.items():
        model = MLP(10, [64,32], 1, activation="relu")

        train_losses, val_losses, _ = train(model, X_train, y_train, X_val, y_val, weight_decay=wd)

        loss_dict[name + " - Train"] = train_losses
        loss_dict[name + " - Val"] = val_losses

    plot(loss_dict, "Experiment D: Regularization Comparison")



#-------------------------------------------------- PLOTTING --------------------------------------------------
def plot(loss_dict, title):
    for label, losses in loss_dict.items():
        plt.plot(losses, label=label)
    plt.xlabel("Epoch")
    plt.ylabel("Loss")
    plt.title(title)
    plt.legend()
    plt.show()



#-------------------------------------------------- MAIN --------------------------------------------------
if __name__ == "__main__":
    experiment_activation_functions()
    experiment_depth_vs_width()
    experiment_optimizers()
    experiment_regularization()