import torch
import torch.nn as nn
import torch.optim as optim
from torch.cuda.amp import GradScaler, autocast
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, confusion_matrix
from sklearn.utils.class_weight import compute_class_weight
from sklearn.model_selection import train_test_split
import numpy as np
import os

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'

# 自定义数据集
class MyDataset(Dataset):
    def __init__(self, specs, labels, augment=False):
        self.specs = specs
        self.labels = labels
        self.augment = augment

    def __getitem__(self, index):
        spec, target = self.specs[index], self.labels[index]
        if self.augment:
            noise = 0.01 * torch.randn_like(spec)
            spec = spec + noise
        return spec, target

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

# 数据标准化
def ZspProcess(X_train, X_test, y_train, y_test, need=True):
    if need:
        scaler = StandardScaler()
        X_train = scaler.fit_transform(X_train)
        X_test = scaler.transform(X_test)

    X_train = torch.tensor(X_train[:, np.newaxis, :], dtype=torch.float32)
    X_test = torch.tensor(X_test[:, np.newaxis, :], dtype=torch.float32)
    y_train = torch.tensor(y_train, dtype=torch.long)
    y_test = torch.tensor(y_test, dtype=torch.long)

    data_train = MyDataset(X_train, y_train, augment=True)
    data_test = MyDataset(X_test, y_test, augment=False)
    return data_train, data_test

# Focal Loss
class FocalLoss(nn.Module):
    def __init__(self, alpha=1, gamma=2, reduction='mean'):
        super(FocalLoss, self).__init__()
        self.alpha = alpha
        self.gamma = gamma
        self.reduction = reduction

    def forward(self, inputs, targets):
        probs = torch.softmax(inputs, dim=1)
        target_probs = probs[range(len(targets)), targets]
        focal_weight = self.alpha * (1 - target_probs) ** self.gamma
        log_prob = -torch.log(target_probs)
        loss = focal_weight * log_prob

        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:
            return loss

# Transformer模块
class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, ff_dim, dropout=0.1):
        super(TransformerBlock, self).__init__()
        self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout)
        self.feed_forward = nn.Sequential(
            nn.Linear(embed_dim, ff_dim),
            nn.ReLU(),
            nn.Linear(ff_dim, embed_dim)
        )
        self.layernorm1 = nn.LayerNorm(embed_dim)
        self.layernorm2 = nn.LayerNorm(embed_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        attn_output, _ = self.attention(x, x, x)
        x = self.layernorm1(x + self.dropout(attn_output))
        ff_output = self.feed_forward(x)
        x = self.layernorm2(x + self.dropout(ff_output))
        return x

# 改进后的CNN+Transformer模型
class CNNWithTransformer(nn.Module):
    def __init__(self, nls, embed_dim=96, num_heads=2, ff_dim=192, dropout=0.1):
        super(CNNWithTransformer, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv1d(1, 64, kernel_size=5, stride=2, padding=2),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(0.2),  # 添加Dropout
            nn.MaxPool1d(2, 2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv1d(64, embed_dim, kernel_size=5, stride=2, padding=2),
            nn.BatchNorm1d(embed_dim),
            nn.ReLU(),
            nn.Dropout(0.2),  # 添加Dropout
            nn.MaxPool1d(2, 2)
        )
        self.transformer = TransformerBlock(embed_dim, num_heads, ff_dim, dropout)
        self.fc = nn.Sequential(
            nn.Linear(embed_dim, 128),
            nn.ReLU(),
            nn.Dropout(0.3),  # 添加Dropout
            nn.Linear(128, nls)
        )

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.permute(2, 0, 1)
        x = self.transformer(x)
        x = x.mean(dim=0)
        x = self.fc(x)
        return x

# 训练函数（包含早停机制）
def TransformerTrain(X_train, X_val, y_train, y_val, BATCH_SIZE, n_epochs, nls, model_path, patience=10):
    data_train, data_val = ZspProcess(X_train, X_val, y_train, y_val, need=True)
    train_loader = DataLoader(data_train, batch_size=BATCH_SIZE, shuffle=True)
    val_loader = DataLoader(data_val, batch_size=BATCH_SIZE, shuffle=False)

    model = CNNWithTransformer(nls=nls).to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=0.001)
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5)
    criterion = FocalLoss(alpha=1, gamma=2).to(device)
    scaler = GradScaler()

    best_val_loss = float('inf')
    early_stop_counter = 0
    y_true_train, y_pred_train = [], []

    for epoch in range(n_epochs):
        model.train()
        train_loss, train_acc = [], []

        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()

            with autocast():
                outputs = model(inputs)
                loss = criterion(outputs, labels)

            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            _, preds = torch.max(outputs, 1)
            y_true_train.extend(labels.cpu().numpy())
            y_pred_train.extend(preds.cpu().numpy())
            acc = accuracy_score(labels.cpu(), preds.cpu())
            train_loss.append(loss.item())
            train_acc.append(acc)

        # 验证集评估
        model.eval()
        val_loss = []
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, labels)
                val_loss.append(loss.item())

        avg_val_loss = np.mean(val_loss)
        avg_train_loss = np.mean(train_loss)
        avg_train_acc = np.mean(train_acc)

        print(f"Epoch [{epoch+1}/{n_epochs}] - Train Loss: {avg_train_loss:.4f}, Train Acc: {avg_train_acc:.4f}, Val Loss: {avg_val_loss:.4f}")

        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            early_stop_counter = 0
            torch.save(model.state_dict(), model_path)
            print("Model improved and saved.")
        else:
            early_stop_counter += 1
            print(f"No improvement. Early stop counter: {early_stop_counter}/{patience}")

        if early_stop_counter >= patience:
            print("Early stopping triggered.")
            break


    # 训练集指标
    train_accuracy = accuracy_score(y_true_train, y_pred_train)
    train_precision = precision_score(y_true_train, y_pred_train, average='weighted')
    train_recall = recall_score(y_true_train, y_pred_train, average='weighted')
    train_f1 = f1_score(y_true_train, y_pred_train, average='weighted')
    train_cm = confusion_matrix(y_true_train, y_pred_train)

    train_metrics = {
        "accuracy": train_accuracy,
        "precision": train_precision,
        "recall": train_recall,
        "f1_score": train_f1,
        "confusion_matrix": train_cm
    }

    return model, train_metrics

# 测试函数
def TransformerTest(X_test, y_test, BATCH_SIZE, nls, model_path):
    data_test = ZspProcess(X_test, X_test, y_test, y_test, need=True)[1]
    test_loader = DataLoader(data_test, batch_size=BATCH_SIZE, shuffle=False)

    model = CNNWithTransformer(nls=nls).to(device)
    model.load_state_dict(torch.load(model_path))
    model.eval()

    y_true, y_pred = [], []
    test_loss = []

    criterion = FocalLoss(alpha=1, gamma=2).to(device)  # 使用 FocalLoss

    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            loss = criterion(outputs, labels)

            _, preds = torch.max(outputs, 1)
            y_true.extend(labels.cpu().numpy())
            y_pred.extend(preds.cpu().numpy())
            test_loss.append(loss.item())

    # 测试集指标
    test_accuracy = accuracy_score(y_true, y_pred)
    test_precision = precision_score(y_true, y_pred, average='weighted')
    test_recall = recall_score(y_true, y_pred, average='weighted')
    test_f1 = f1_score(y_true, y_pred, average='weighted')
    test_cm = confusion_matrix(y_true, y_pred)

    test_metrics = {
        "accuracy": test_accuracy,
        "precision": test_precision,
        "recall": test_recall,
        "f1_score": test_f1,
        "confusion_matrix": test_cm
    }

    print(f"Accuracy: {test_accuracy:.4f}, Precision: {test_precision:.4f}, Recall: {test_recall:.4f}, F1 Score: {test_f1:.4f}")
    print(f"Confusion Matrix:\n{test_cm}")
    return test_metrics


def TransformerTrainAndTest(X,X_test,  y, y_test, BATCH_SIZE, n_epochs, nls, model_path, val_split=0.2, patience=10):
    # 从训练集中划分验证集
    X_train, X_val, y_train, y_val = train_test_split(X, y, test_size=val_split, random_state=42)

    # 训练模型并获取训练指标
    model, train_metrics = TransformerTrain(X_train, X_val, y_train, y_val, BATCH_SIZE, n_epochs, nls, model_path, patience)

    # 测试模型并获取测试指标
    test_metrics = TransformerTest(X_test, y_test, BATCH_SIZE, nls, model_path)

    return train_metrics, test_metrics