import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
import numpy as np


# Define a simple dataset
class RepDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        return self.data[idx], self.labels[idx]


class MLPClassifier(nn.Module):
    def __init__(self, input_dim, num_classes):
        super(MLPClassifier, self).__init__()
        self.fc = nn.Linear(input_dim, num_classes)

    def forward(self, x):
        return self.fc(x)


# Prediction function
def predict(model, dataloader, device):
    model.eval()  # Set the model to evaluation mode
    all_predictions = []
    total = 0
    correct = 0
    with torch.no_grad():  # Disable gradient calculation
        for batch_data, batch_labels in dataloader:
            batch_data, batch_labels = batch_data.to(device), batch_labels.to(device)
            outputs = model(batch_data)
            _, predicted = torch.max(outputs, 1)
            all_predictions.extend(predicted.cpu().numpy())
            total += batch_labels.size(0)
            correct += (predicted == batch_labels).sum().item()
    accuracy = correct / total
    return all_predictions, accuracy


def train_predict_mlp(
    positive_train_states,
    positive_test_states,
    negative_train_states,
    negative_test_states,
    classifier_save_name="model",
):

    np.random.seed(0)
    torch.manual_seed(0)

    X_train = torch.cat((positive_train_states, negative_train_states), dim=0)

    y_train = [1 for _ in range(positive_train_states.size(0))] + [
        0 for _ in range(negative_train_states.size(0))
    ]

    X_test = torch.cat((positive_test_states, negative_test_states), dim=0)

    y_test = [1 for _ in range(positive_test_states.size(0))] + [
        0 for _ in range(negative_test_states.size(0))
    ]

    # Create dataset and dataloader
    dataset_train = RepDataset(X_train, y_train)
    dataloader_train = DataLoader(dataset_train, batch_size=128, shuffle=True)

    dataset_test = RepDataset(X_test, y_test)
    dataloader_test = DataLoader(dataset_test, batch_size=128, shuffle=True)

    # Initialize the model, loss function, and optimizer

    model = MLPClassifier(X_train.size(1), 2)

    # Check if GPU is available and move the model to GPU if possible
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(model.parameters(), lr=0.008)

    # Training loop
    num_epochs = 300
    for epoch in range(num_epochs):
        for batch_data, batch_labels in dataloader_train:
            # Forward pass
            batch_data, batch_labels = batch_data.to(device), batch_labels.to(device)
            outputs = model(batch_data)
            loss = criterion(outputs, batch_labels)

            # Backward pass and optimization
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")
    print("Training complete.")

    y_pred, test_acc = predict(model, dataloader_test, device)

    # Save the trained model
    model_path = classifier_save_name + ".pth"
    torch.save(model.state_dict(), model_path)
    print(f"Model saved to {model_path}")

    print(f"Accuracy: {test_acc}")

    return y_pred
