import pandas as pd
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torch.optim import Adam
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import roc_curve, auc, accuracy_score, precision_score, recall_score, f1_score
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import numpy as np
import os


import matplotlib
matplotlib.rcParams['font.family'] = 'Times New Roman'
matplotlib.rcParams['savefig.dpi'] = 300
matplotlib.rcParams['font.size'] = 24
matplotlib.rcParams['axes.labelweight'] = 'bold'
matplotlib.rcParams['axes.titleweight'] = 'bold'
matplotlib.rcParams['font.weight'] = 'bold'

class CollapseDataset(Dataset):
    def __init__(self, features, labels):
        self.features = features
        self.labels = labels

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

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

def load_data(file_path):
    df = pd.read_excel(file_path)
    X = df.iloc[:, :9].values  
    y = df.iloc[:, 9].values  
    return X, y

class PrototypicalNetwork(nn.Module):
    def __init__(self, input_size, hidden_size):
        super(PrototypicalNetwork, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size // 2),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(hidden_size // 2, hidden_size // 4),
            nn.ReLU()
        )

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

def train(model, dataloader, optimizer, criterion):
    model.train()
    total_loss = 0
    for data, target in dataloader:
        optimizer.zero_grad()
        output = model(data.float())
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        total_loss += loss.item()
    return total_loss / len(dataloader)

def evaluate(model, dataloader):
    model.eval()
    y_true = []
    y_scores = []
    y_pred = []
    embeddings = []
    softmax = torch.nn.Softmax(dim=1)
    with torch.no_grad():
        for data, target in dataloader:
            output = model(data.float())
            embeddings.append(output)
            scores = softmax(output)[:, 1]  # Assuming index 1 is the positive class
            _, predicted = torch.max(output.data, 1)
            y_true.extend(target.tolist())
            y_scores.extend(scores.tolist())
            y_pred.extend(predicted.tolist())
    embeddings = torch.cat(embeddings).cpu().numpy()

    accuracy = accuracy_score(y_true, y_pred)
    precision = precision_score(y_true, y_pred)
    recall = recall_score(y_true, y_pred)
    f1 = f1_score(y_true, y_pred)

    return y_true, y_scores, embeddings, accuracy, precision, recall, f1

def plot_and_save_roc_curve(y_true, y_scores, save_path):
    fpr, tpr, _ = roc_curve(y_true, y_scores)
    roc_auc = auc(fpr, tpr)
    plt.figure(figsize=(10, 8))
    lw = 2
    plt.plot(fpr, tpr, color='darkorange', lw=lw, label=f'ROC curve (area = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=lw, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('Receiver Operating Characteristic')
    plt.legend(loc="lower right")
    plt.savefig(save_path)
    plt.close()

def visualize_and_save_embeddings_2d(embeddings, targets, title, save_path):
    tsne = TSNE(n_components=2, random_state=42)
    transformed_embeddings = tsne.fit_transform(embeddings)
    plt.figure(figsize=(10, 8))
    colors = ['red', 'green', 'blue', 'orange', 'purple']
    for i, target in enumerate(np.unique(targets)):
        indices = np.where(targets == target)
        plt.scatter(transformed_embeddings[indices, 0], transformed_embeddings[indices, 1],
                    s=100, c=colors[i % len(colors)], label=f'Class {target}')
    plt.title(title)
    plt.legend()
    plt.savefig(save_path)
    plt.close()

def save_predictions(save_path, y_true, y_pred, y_scores):
    df = pd.DataFrame({
        'True Labels': y_true,
        'Predicted Labels': y_pred,
        'Scores': y_scores
    })
    df.to_csv(save_path, index=False)

def plot_and_save_feature_importance(model, feature_names, save_path):
    # Assuming we use the weights of the first linear layer as the measure of feature importance
    importance = model.encoder[0].weight.detach().cpu().numpy().mean(axis=0)
    sorted_idx = np.argsort(importance)

    plt.figure(figsize=(12, 8))
    plt.barh(range(len(sorted_idx)), importance[sorted_idx], align='center')
    plt.yticks(range(len(sorted_idx)), [feature_names[i] for i in sorted_idx])
    plt.xlabel('Importance')
    plt.title('Feature Importance')
    plt.tight_layout()
    plt.savefig(save_path)
    plt.close()

def main():
    train_file_path = r""
    test_file_path = r""
    output_dir = r""
    os.makedirs(output_dir, exist_ok=True)

    # Load and standardize data
    X_train, y_train = load_data(train_file_path)
    X_test, y_test = load_data(test_file_path)
    scaler = StandardScaler()
    X_train_scaled = scaler.fit_transform(X_train)
    X_test_scaled = scaler.transform(X_test)

    # Data loading
    train_dataset = CollapseDataset(torch.tensor(X_train_scaled), torch.tensor(y_train, dtype=torch.long))
    test_dataset = CollapseDataset(torch.tensor(X_test_scaled), torch.tensor(y_test, dtype=torch.long))
    train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
    test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False)

    # Model and optimizer initialization
    input_size = 9
    hidden_size = 64
    model = PrototypicalNetwork(input_size, hidden_size)
    optimizer = Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()

    # Training
    for epoch in range(10):
        train_loss = train(model, train_loader, optimizer, criterion)
        print(f'Epoch {epoch + 1}, Loss: {train_loss:.4f}')

    # Evaluation and result saving
    y_true, y_scores, embeddings, accuracy, precision, recall, f1 = evaluate(model, test_loader)
    y_pred = [int(score >= 0.5) for score in y_scores]  # Converting scores to binary predictions
    save_predictions(f"{output_dir}/test_predictions.csv", y_true, y_pred, y_scores)
    plot_and_save_roc_curve(y_true, y_scores, f"{output_dir}/test_roc_curve.png")
    visualize_and_save_embeddings_2d(embeddings, y_true, 't-SNE Embeddings (Test Set)', f"{output_dir}/test_tsne_2d.png")
    feature_names = [f'Feature {i}' for i in range(input_size)]
    plot_and_save_feature_importance(model, feature_names, f"{output_dir}/feature_importance.png")

    print(f"Accuracy: {accuracy:.4f}")
    print(f"Precision: {precision:.4f}")
    print(f"Recall: {recall:.4f}")
    print(f"F1 Score: {f1:.4f}")

if __name__ == "__main__":
    main()
