import os
import time
from wsgiref import validate
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import Dataset, DataLoader, TensorDataset
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score, confusion_matrix, roc_curve, precision_recall_curve, average_precision_score
from sklearn.model_selection import train_test_split
import json
import pickle
from imblearn.over_sampling import SMOTE
from collections import defaultdict
import lightgbm as lgb
import optuna
from optuna import Trial
from optuna.samplers import TPESampler
import matplotlib.pyplot as plt
import seaborn as sns

# 记录开始时间
start_time = time.time()

BATCH_SIZE = 512
EPOCHS = 100
LEARNING_RATE = 1e-5
DATA_SPLIT = [0.8, 0.1, 0.1]
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
NUM_FOLDS = 10  
ENABLE_DOWNSAMPLING = True  
STRATIFIED_BINS = 45  
OPTUNA_N_TRIALS = 50  # Optuna试验次数

def set_seed(seed=42):
    import random
    import os
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)  
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

    # 新增环境变量防止并行操作引入随机性
    os.environ['OMP_NUM_THREADS'] = '1'
    os.environ['MKL_NUM_THREADS'] = '1'
    torch.set_num_threads(1)  # 限制PyTorch线程数

set_seed(42)  

def load_labels(label_path):
    df = pd.read_csv(label_path, sep='\t')
    positive_samples = set()
    for _, row in df.iterrows():
        # 确保统一使用大写形式
        key = (
            row['pdbID'].strip().upper(),
            str(row['chain']).strip().upper(),
            str(row['res_num']).strip()
        )
        positive_samples.add(key)
    return positive_samples


class ResidueDataset(Dataset):
    def __init__(self, esm_feature_dir, label_csv, pdb_id_list=None, window_size=2):
        self.positive_samples = load_labels(label_csv)
        self.samples = []
        self.pdb_samples = {}
        self.window_size = window_size
        
        # 修复点1: 转换输入的PDB ID列表为大写
        if pdb_id_list is not None:
            pdb_id_list = [pdb.upper().strip() for pdb in pdb_id_list]
        
        # 收集可用的特征文件
        feature_files = []
        for pdb_file in os.listdir(esm_feature_dir):
            # 修复点2: 正确的文件名分割方式
            file_parts = pdb_file.split('_features.txt')
            if file_parts and file_parts[0]:
                pdb_id = file_parts[0].upper().strip()
                # 修复点3: 只处理存在于指定列表的PDB
                if pdb_id_list is None or pdb_id in pdb_id_list:
                    feature_files.append((pdb_id, os.path.join(esm_feature_dir, pdb_file)))
        
        print(f"找到 {len(feature_files)} 个特征文件")
        
        for pdb_id, feature_path in feature_files:
            pdb_data = []
            try:
                with open(feature_path, 'r') as f:
                    for line in f:
                        parts = line.strip().split('\t')
                        if len(parts) < 5:
                            continue
                        
                        # 修复点4: 统一使用大写
                        file_pdb_id = parts[0].upper().strip()
                        chain = parts[1].upper().strip()
                        res_num = parts[2].strip()
                        
                        # 确保特征长度正确
                        features = np.array(parts[4].split(), dtype=np.float32)
                        if len(features) != 1024:
                            continue
                        
                        pdb_data.append({
                            'pdb_id': file_pdb_id,
                            'chain': chain,
                            'res_num': res_num,
                            'features': features
                        })
            except Exception as e:
                print(f"读取文件 {feature_path} 时出错: {str(e)}")
                continue
            
            if not pdb_data:
                print(f"警告: PDB {pdb_id} 没有有效数据")
                continue
            
            # 按链分组
            chain_dict = defaultdict(list)
            for data in pdb_data:
                chain_dict[data['chain']].append(data)
            
            # 为每条链计算上下文特征
            for chain, residues in chain_dict.items():
                residues = sorted(residues, key=lambda x: int(x['res_num']))
                
                # 计算上下文特征
                for i in range(len(residues)):
                    # 确定窗口范围
                    start_idx = max(0, i - self.window_size)
                    end_idx = min(len(residues), i + self.window_size + 1)
                    
                    # 收集窗口内残基的特征（排除当前残基）
                    context_features = []
                    for j in range(start_idx, end_idx):
                        if j != i:  # 排除自身
                            context_features.append(residues[j]['features'])
                    
                    # 计算上下文统计量
                    if context_features:
                        mean_context = np.mean(context_features, axis=0)
                        std_context = np.std(context_features, axis=0)
                    else:
                        # 如果没有上下文，用零填充
                        mean_context = np.zeros(1024)
                        std_context = np.zeros(1024)
                    
                    # 创建增强特征向量
                    extended_features = np.concatenate([
                        residues[i]['features'],  # 原始特征
                        mean_context,              # 上下文均值
                        std_context                # 上下文标准差
                    ])
                    
                    # 创建样本key并确定标签
                    key = (residues[i]['pdb_id'], residues[i]['chain'], residues[i]['res_num'])
                    label = 1 if key in self.positive_samples else 0
                    
                    # 添加到样本列表
                    self.samples.append({
                        'features': extended_features,
                        'label': label,
                        'pdb_id': residues[i]['pdb_id'],
                        'chain': residues[i]['chain'],
                        'res_num': residues[i]['res_num']
                    })
            
            # 为数据集分割功能记录PDB索引
            if pdb_id not in self.pdb_samples:
                self.pdb_samples[pdb_id] = []
            for i in range(len(self.samples) - len(residues), len(self.samples)):
                self.pdb_samples[pdb_id].append(i)
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        item = self.samples[idx]
        return (
            torch.tensor(item['features'], dtype=torch.float32),
            torch.tensor(item['label'], dtype=torch.float32)
        )

def split_dataset_by_pdb(dataset, pdb_list):
    indices = []
    # 将输入的pdb_list转换为大写
    pdb_list = [pdb.upper().strip() for pdb in pdb_list]
    
    # 确保只使用数据集中存在的PDB
    for pdb in pdb_list:
        if pdb in dataset.pdb_samples:
            indices.extend(dataset.pdb_samples[pdb])
        else:
            print(f"警告：PDB {pdb} 不在数据集中，已跳过")
    return indices

def get_stratified_fold_indices(dataset):
    # 计算每个PDB的正样本个数
    pdb_positive_counts = {}
    for pdb, indices in dataset.pdb_samples.items():
        pos_count = sum(1 for idx in indices if dataset.samples[idx]['label'] == 1)
        pdb_positive_counts[pdb] = pos_count
    
    # 将PDB按照正样本个数分组（分桶）
    bins = STRATIFIED_BINS
    bins_dict = {i: [] for i in range(bins)}
    
    # 确定每个桶的正样本个数范围
    min_count = min(pdb_positive_counts.values())
    max_count = max(pdb_positive_counts.values())
    bin_width = (max_count - min_count) / bins
    
    # 将PDB分配到对应的桶中
    for pdb, count in pdb_positive_counts.items():
        # 计算PDB属于哪个桶
        bin_index = min(int((count - min_count) / bin_width), bins - 1)
        bins_dict[bin_index].append(pdb)
    
    # 初始化折的列表
    fold_indices = [[] for _ in range(NUM_FOLDS)]
    
    # 对每个桶内的PDB进行分层抽样
    for bin_index, pdbs_in_bin in bins_dict.items():
        # 打乱当前桶内的PDB列表
        np.random.shuffle(pdbs_in_bin)
        
        # 将当前桶内的PDB分配到各折中
        for idx, pdb in enumerate(pdbs_in_bin):
            fold_idx = idx % NUM_FOLDS
            fold_indices[fold_idx].append(pdb)
    
    return fold_indices

# ==== 修改后的MLP+RES模型 ====
class ResidualBlock(nn.Module):
    def __init__(self, input_dim, hidden_dim, dropout_rate=0.1):
        super(ResidualBlock, self).__init__()
        self.linear1 = nn.Linear(input_dim, hidden_dim)
        self.bn1 = nn.BatchNorm1d(hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, input_dim)
        self.bn2 = nn.BatchNorm1d(input_dim)
        self.dropout = nn.Dropout(dropout_rate)
        self.activation = nn.ReLU()
        
        # 如果输入输出维度不匹配，使用1x1卷积调整维度
        if input_dim != hidden_dim:
            self.shortcut = nn.Sequential(
                nn.Linear(input_dim, input_dim),
                nn.BatchNorm1d(input_dim)
            )
        else:
            self.shortcut = nn.Identity()

    def forward(self, x):
        identity = self.shortcut(x)
        
        out = self.linear1(x)
        out = self.bn1(out)
        out = self.activation(out)
        out = self.dropout(out)
        
        out = self.linear2(out)
        out = self.bn2(out)
        
        out += identity
        out = self.activation(out)
        out = self.dropout(out)
        
        return out

class MLPWithResidual(nn.Module):
    def __init__(self, input_dim, hidden_dims, num_res_blocks, dropout_rate=0.1):
        super(MLPWithResidual, self).__init__()
        
        self.input_layer = nn.Sequential(
            nn.Linear(input_dim, hidden_dims[0]),
            nn.BatchNorm1d(hidden_dims[0]),
            nn.ReLU(),
            nn.Dropout(dropout_rate)
        )
        
        # 创建残差块 - 修改为1-3个
        self.res_blocks = nn.ModuleList()
        current_dim = hidden_dims[0]
        
        for i in range(num_res_blocks):
            self.res_blocks.append(ResidualBlock(current_dim, hidden_dims[min(1, len(hidden_dims)-1)], dropout_rate))
        
        # 隐藏层 - 修改为1-3层
        self.hidden_layers = nn.ModuleList()
        for i in range(1, len(hidden_dims)):
            self.hidden_layers.append(nn.Sequential(
                nn.Linear(hidden_dims[i-1], hidden_dims[i]),
                nn.BatchNorm1d(hidden_dims[i]),
                nn.ReLU(),
                nn.Dropout(dropout_rate)
            ))
        
        self.output_layer = nn.Linear(hidden_dims[-1], 1)
        
    def forward(self, x):
        x = self.input_layer(x)
        
        # 通过残差块
        for res_block in self.res_blocks:
            x = res_block(x)
        
        # 通过隐藏层
        for layer in self.hidden_layers:
            x = layer(x)
            
        x = self.output_layer(x)
        return x

class ContactPredictor:
    def __init__(self, params=None):
        if params is None:
            params = {
                'hidden_dims': [512, 256, 128],
                'num_res_blocks': 2,
                'dropout_rate': 0.2,
                'learning_rate': 1e-4,
                'weight_decay': 1e-5
            }
        
        self.model = None
        self.params = params
        self.device = DEVICE
        
    def fit(self, train_loader, val_loader):
        """训练MLP+RES模型"""
        input_dim = train_loader.dataset[0][0].shape[0]
        
        # 修复点：检查参数中是否包含隐藏层维度信息
        if 'hidden_dims' in self.params:
            # 如果直接提供了hidden_dims，使用它
            hidden_dims = self.params['hidden_dims']
        elif 'hidden_dim_base' in self.params:
            # 如果使用Optuna的参数，动态构建hidden_dims
            hidden_dims = []
            current_dim = self.params['hidden_dim_base']
            # 修改：隐藏层层数范围改为1-3
            num_layers = min(3, max(1, self.params['hidden_layers']))
            for i in range(num_layers):
                hidden_dims.append(int(current_dim))
                current_dim = current_dim * self.params['reduction_factor']
        else:
            # 默认隐藏层维度
            hidden_dims = [512, 256, 128]
        
        # 确保hidden_dims是列表格式且长度在1-3之间
        if isinstance(hidden_dims, (int, float)):
            hidden_dims = [int(hidden_dims)]
        # 限制隐藏层层数为1-3
        if len(hidden_dims) > 3:
            hidden_dims = hidden_dims[:3]
        
        # 获取其他参数，提供默认值，并限制残差块个数为1-3
        num_res_blocks = min(3, max(1, self.params.get('num_res_blocks', 2)))
        dropout_rate = self.params.get('dropout_rate', 0.2)
        learning_rate = self.params.get('learning_rate', 1e-4)
        weight_decay = self.params.get('weight_decay', 1e-5)
        
        print(f"训练参数: hidden_dims={hidden_dims}, num_res_blocks={num_res_blocks}, "
              f"dropout_rate={dropout_rate:.4f}, lr={learning_rate:.2e}, weight_decay={weight_decay:.2e}")
        
        self.model = MLPWithResidual(
            input_dim=input_dim,
            hidden_dims=hidden_dims,
            num_res_blocks=num_res_blocks,
            dropout_rate=dropout_rate
        ).to(self.device)
        
        criterion = FocalLoss()
        optimizer = torch.optim.AdamW(
            self.model.parameters(),
            lr=learning_rate,
            weight_decay=weight_decay
        )
        
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer, mode='max', factor=0.5, patience=10, verbose=True
        )
        
        best_val_auc = 0
        patience_counter = 0
        patience = 20
        
        for epoch in range(EPOCHS):
            # 训练阶段
            self.model.train()
            train_loss = 0
            for inputs, labels in train_loader:
                inputs, labels = inputs.to(self.device), labels.to(self.device)
                
                optimizer.zero_grad()
                outputs = self.model(inputs).squeeze()
                loss = criterion(outputs, labels)
                loss.backward()
                optimizer.step()
                
                train_loss += loss.item()
            
            # 验证阶段
            val_auc, top10_acc = self.evaluate(val_loader)
            scheduler.step(val_auc)
            
            if (epoch + 1) % 10 == 0:
                print(f'Epoch {epoch+1}/{EPOCHS}, Train Loss: {train_loss/len(train_loader):.4f}, '
                      f'Val AUC: {val_auc:.4f}, Top10 Acc: {top10_acc:.2%}')
            
            # 早停
            if val_auc > best_val_auc:
                best_val_auc = val_auc
                patience_counter = 0
                # 保存最佳模型
                torch.save(self.model.state_dict(), 'best_model.pth')
            else:
                patience_counter += 1
                if patience_counter >= patience:
                    print(f'Early stopping at epoch {epoch+1}')
                    break
        
        # 加载最佳模型
        if os.path.exists('best_model.pth'):
            self.model.load_state_dict(torch.load('best_model.pth'))
        print(f'训练完成 | Best Val AUC: {best_val_auc:.4f}')
    
    def evaluate(self, data_loader):
        """评估模型"""
        if self.model is None:
            return 0.0, 0.0
            
        self.model.eval()
        all_probs = []
        all_labels = []
        
        with torch.no_grad():
            for inputs, labels in data_loader:
                inputs, labels = inputs.to(self.device), labels.to(self.device)
                outputs = self.model(inputs).squeeze()
                probs = torch.sigmoid(outputs)
                all_probs.extend(probs.cpu().numpy())
                all_labels.extend(labels.cpu().numpy())
        
        all_probs = np.array(all_probs)
        all_labels = np.array(all_labels)
        
        # 计算AUC和top10准确率
        try:
            # 修复：确保有足够的样本计算top10
            if len(all_probs) < 500:
                topk = len(all_probs)
            else:
                topk = 500
                
            sorted_idx = np.argsort(all_probs)[::-1]
            top10_labels = all_labels[sorted_idx[:topk]]
            top10_acc = np.mean(top10_labels)
            
            auc = roc_auc_score(all_labels, all_probs)
            return auc, top10_acc
        except Exception as e:
            print(f"评估过程中出错: {str(e)}")
            return 0.0, 0.0
    
    def predict_proba(self, X):
        """预测概率"""
        if self.model is None:
            raise ValueError("模型尚未训练，请先调用fit方法")
            
        self.model.eval()
        if isinstance(X, torch.Tensor):
            X_tensor = X.to(self.device)
        else:
            X_tensor = torch.tensor(X, dtype=torch.float32).to(self.device)
        
        with torch.no_grad():
            outputs = self.model(X_tensor).squeeze()
            probs = torch.sigmoid(outputs)
        
        return probs.cpu().numpy()
    
    def predict(self, X):
        """预测类别"""
        probs = self.predict_proba(X)
        return np.round(probs)

# ==== 修改后的Optuna目标函数 ====
def objective(trial, train_loader, val_loader):
    """Optuna优化目标函数 - MLP+RES版本"""
    
    # 定义MLP+RES的超参数搜索空间 - 修改为1-3层和1-3个残差块
    params = {
        # 隐藏层维度选择 - 修改为1-3层
        'hidden_dim_base': trial.suggest_int('hidden_dim_base', 256, 1024),
        'hidden_layers': trial.suggest_int('hidden_layers', 1, 3),  # 改为1-3
        'reduction_factor': trial.suggest_float('reduction_factor', 0.3, 0.7),
        
        # 残差块数量 - 修改为1-3个
        'num_res_blocks': trial.suggest_int('num_res_blocks', 1, 3),  # 改为1-3
        
        # 正则化参数
        'dropout_rate': trial.suggest_float('dropout_rate', 0.1, 0.5),
        'learning_rate': trial.suggest_float('learning_rate', 1e-5, 1e-3, log=True),
        'weight_decay': trial.suggest_float('weight_decay', 1e-6, 1e-4, log=True),
    }
    
    # 构建隐藏层维度列表
    hidden_dims = []
    current_dim = params['hidden_dim_base']
    # 确保层数在1-3范围内
    num_layers = min(3, max(1, params['hidden_layers']))
    for i in range(num_layers):
        hidden_dims.append(int(current_dim))
        current_dim = current_dim * params['reduction_factor']
    
    params['hidden_dims'] = hidden_dims
    
    try:
        # 创建并训练模型
        model = ContactPredictor(params)
        model.fit(train_loader, val_loader)
        
        # 在验证集上评估
        val_auc, _ = model.evaluate(val_loader)
        
        # 清理模型文件
        if os.path.exists('best_model.pth'):
            os.remove('best_model.pth')
            
        return val_auc
    except Exception as e:
        # 如果训练失败，返回一个很低的分数
        print(f"训练失败，参数: {params}, 错误: {str(e)}")
        # 清理可能创建的模型文件
        if os.path.exists('best_model.pth'):
            os.remove('best_model.pth')
        return 0.0

# ==== FocalLoss类保持不变 ====
class FocalLoss(nn.Module):
    def __init__(self, alpha=0.8, gamma=1.5):
        super().__init__()
        self.alpha = alpha
        self.gamma = gamma

    def forward(self, inputs, targets):
        BCE_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
        pt = torch.exp(-BCE_loss)
        focal_loss = self.alpha * (1 - pt) ** self.gamma * BCE_loss
        return focal_loss.mean()

def evaluate_with_pdb_output(dataset, test_indices, pred_probs, output_dir="test_results"):
    os.makedirs(output_dir, exist_ok=True)
    pdb_results = {}
    
    for idx, prob in zip(test_indices, pred_probs):
        data = dataset.samples[idx]
        pdb_id = data['pdb_id']
        record = {
            'chain': data['chain'],
            'res_num': data['res_num'],
            'probability': prob,
            'label': data['label']
        }
        if pdb_id not in pdb_results:
            pdb_results[pdb_id] = []
        pdb_results[pdb_id].append(record)
    
    K_LIST = [5, 10, 20, 50, 100]
    for pdb_id, records in pdb_results.items():
        sorted_records = sorted(records, key=lambda x: x['probability'], reverse=True)
        
        with open(os.path.join(output_dir, f"{pdb_id}_results.csv"), 'w') as f:
            f.write("Chain,ResNum,Probability,Label\n")
            for r in sorted_records:
                f.write(f"{r['chain']},{r['res_num']},{r['probability']:.4f},{r['label']}\n")
            
            f.write("\n\n=== TopK accuracy ===\n")
            for k in K_LIST:
                topk = sorted_records[:k]
                total = len(topk)
                if total == 0:
                    continue
                
                correct = sum(1 for r in topk if r['label'] == 1)
                acc = correct / total
                residues = [(r['chain'], r['res_num']) for r in topk]
                
                f.write(
                    f"Top{k}accuracy: {acc:.2%} ({correct}/{total})\n"
                    f"residues: {', '.join([f'{c}{r}' for c, r in residues])}\n\n"
                )

def calculate_top1_accuracy(dataset, test_indices, probs):
    pdb_results = {}
    
    for idx, prob in zip(test_indices, probs):
        data = dataset.samples[idx]
        pdb_id = data['pdb_id']
        if pdb_id not in pdb_results:
            pdb_results[pdb_id] = []
        pdb_results[pdb_id].append({
            'chain': data['chain'],
            'res_num': data['res_num'],
            'probability': prob,
            'label': data['label']
        })
    
    correct_count = 0
    total_pdbs = len(pdb_results)
    
    for pdb_id, records in pdb_results.items():
        sorted_records = sorted(records, key=lambda x: x['probability'], reverse=True)
        top1 = sorted_records[:1]
        if any(r['label'] == 1 for r in top1):
            correct_count += 1
    
    return correct_count / total_pdbs if total_pdbs > 0 else 0

def calculate_top5_accuracy(dataset, test_indices, probs):
    pdb_results = {}
    
    for idx, prob in zip(test_indices, probs):
        data = dataset.samples[idx]
        pdb_id = data['pdb_id']
        if pdb_id not in pdb_results:
            pdb_results[pdb_id] = []
        pdb_results[pdb_id].append({
            'chain': data['chain'],
            'res_num': data['res_num'],
            'probability': prob,
            'label': data['label']
        })
    
    correct_count = 0
    total_pdbs = len(pdb_results)
    
    for pdb_id, records in pdb_results.items():
        sorted_records = sorted(records, key=lambda x: x['probability'], reverse=True)
        top5 = sorted_records[:5]
        if any(r['label'] == 1 for r in top5):
            correct_count += 1
    
    return correct_count / total_pdbs if total_pdbs > 0 else 0

def calculate_top10_accuracy(dataset, test_indices, probs):
    pdb_results = {}
    
    for idx, prob in zip(test_indices, probs):
        data = dataset.samples[idx]
        pdb_id = data['pdb_id']
        if pdb_id not in pdb_results:
            pdb_results[pdb_id] = []
        pdb_results[pdb_id].append({
            'chain': data['chain'],
            'res_num': data['res_num'],
            'probability': prob,
            'label': data['label']
        })
    
    correct_count = 0
    total_pdbs = len(pdb_results)
    
    for pdb_id, records in pdb_results.items():
        sorted_records = sorted(records, key=lambda x: x['probability'], reverse=True)
        top10 = sorted_records[:10]
        if any(r['label'] == 1 for r in top10):
            correct_count += 1
    
    return correct_count / total_pdbs if total_pdbs > 0 else 0

def calculate_top20_accuracy(dataset, test_indices, probs):
    pdb_results = {}
    
    for idx, prob in zip(test_indices, probs):
        data = dataset.samples[idx]
        pdb_id = data['pdb_id']
        if pdb_id not in pdb_results:
            pdb_results[pdb_id] = []
        pdb_results[pdb_id].append({
            'chain': data['chain'],
            'res_num': data['res_num'],
            'probability': prob,
            'label': data['label']
        })
    
    correct_count = 0
    total_pdbs = len(pdb_results)
    
    for pdb_id, records in pdb_results.items():
        sorted_records = sorted(records, key=lambda x: x['probability'], reverse=True)
        top20 = sorted_records[:20]
        if any(r['label'] == 1 for r in top20):
            correct_count += 1
    
    return correct_count / total_pdbs if total_pdbs > 0 else 0

def calculate_top50_accuracy(dataset, test_indices, probs):
    pdb_results = {}
    
    for idx, prob in zip(test_indices, probs):
        data = dataset.samples[idx]
        pdb_id = data['pdb_id']
        if pdb_id not in pdb_results:
            pdb_results[pdb_id] = []
        pdb_results[pdb_id].append({
            'chain': data['chain'],
            'res_num': data['res_num'],
            'probability': prob,
            'label': data['label']
        })
    
    correct_count = 0
    total_pdbs = len(pdb_results)
    
    for pdb_id, records in pdb_results.items():
        sorted_records = sorted(records, key=lambda x: x['probability'], reverse=True)
        top50 = sorted_records[:50]
        if any(r['label'] == 1 for r in top50):
            correct_count += 1
    
    return correct_count / total_pdbs if total_pdbs > 0 else 0

def calculate_top100_accuracy(dataset, test_indices, probs):
    pdb_results = {}
    
    for idx, prob in zip(test_indices, probs):
        data = dataset.samples[idx]
        pdb_id = data['pdb_id']
        if pdb_id not in pdb_results:
            pdb_results[pdb_id] = []
        pdb_results[pdb_id].append({
            'chain': data['chain'],
            'res_num': data['res_num'],
            'probability': prob,
            'label': data['label']
        })
    
    correct_count = 0
    total_pdbs = len(pdb_results)
    
    for pdb_id, records in pdb_results.items():
        sorted_records = sorted(records, key=lambda x: x['probability'], reverse=True)
        top50 = sorted_records[:100]
        if any(r['label'] == 1 for r in top50):
            correct_count += 1
    
    return correct_count / total_pdbs if total_pdbs > 0 else 0

# ==== 新增函数：绘制和保存评估图表 ====
def plot_and_save_evaluation_plots(y_true, y_probs, output_dir, fold_num):
    """绘制并保存AUC曲线、AP曲线和混淆矩阵"""
    
    # 设置matplotlib中文字体和风格
    plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False
    sns.set_style("whitegrid")
    
    # 创建输出目录
    plots_dir = os.path.join(output_dir, "evaluation_plots")
    os.makedirs(plots_dir, exist_ok=True)
    
    # 1. 绘制AUC曲线
    fpr, tpr, thresholds = roc_curve(y_true, y_probs)
    roc_auc = roc_auc_score(y_true, y_probs)
    
    plt.figure(figsize=(8, 6))
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {roc_auc:.4f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='Random')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title(f'Fold {fold_num} - ROC Curve')
    plt.legend(loc="lower right")
    plt.tight_layout()
    plt.savefig(os.path.join(plots_dir, f'fold_{fold_num}_roc_curve.png'), dpi=300, bbox_inches='tight')
    plt.savefig(os.path.join(plots_dir, f'fold_{fold_num}_roc_curve.pdf'), bbox_inches='tight')
    plt.close()
    
    # 保存AUC曲线数据
    auc_data = {
        'fpr': fpr.tolist(),
        'tpr': tpr.tolist(),
        'thresholds': thresholds.tolist(),
        'auc': float(roc_auc)
    }
    with open(os.path.join(plots_dir, f'fold_{fold_num}_roc_data.json'), 'w') as f:
        json.dump(auc_data, f, indent=4)
    
    # 2. 绘制PR曲线
    precision, recall, _ = precision_recall_curve(y_true, y_probs)
    average_precision = average_precision_score(y_true, y_probs)
    
    plt.figure(figsize=(8, 6))
    plt.plot(recall, precision, color='blue', lw=2, 
             label=f'PR curve (AP = {average_precision:.4f})')
    
    # 计算随机分类器的性能
    positive_ratio = np.sum(y_true) / len(y_true)
    plt.axhline(y=positive_ratio, color='red', linestyle='--', 
                label=f'Random (AP = {positive_ratio:.4f})')
    
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('Recall')
    plt.ylabel('Precision')
    plt.title(f'Fold {fold_num} - Precision-Recall Curve')
    plt.legend(loc="upper right")
    plt.tight_layout()
    plt.savefig(os.path.join(plots_dir, f'fold_{fold_num}_pr_curve.png'), dpi=300, bbox_inches='tight')
    plt.savefig(os.path.join(plots_dir, f'fold_{fold_num}_pr_curve.pdf'), bbox_inches='tight')
    plt.close()
    
    # 保存PR曲线数据
    pr_data = {
        'precision': precision.tolist(),
        'recall': recall.tolist(),
        'average_precision': float(average_precision),
        'positive_ratio': float(positive_ratio)
    }
    with open(os.path.join(plots_dir, f'fold_{fold_num}_pr_data.json'), 'w') as f:
        json.dump(pr_data, f, indent=4)
    
    # 3. 绘制混淆矩阵
    y_pred = (y_probs >= 0.5).astype(int)
    cm = confusion_matrix(y_true, y_pred)
    
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=['Predicted Negative', 'Predicted Positive'],
                yticklabels=['Actual Negative', 'Actual Positive'])
    plt.title(f'Fold {fold_num} - Confusion Matrix\n(Threshold = 0.5)')
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.tight_layout()
    plt.savefig(os.path.join(plots_dir, f'fold_{fold_num}_confusion_matrix.png'), dpi=300, bbox_inches='tight')
    plt.savefig(os.path.join(plots_dir, f'fold_{fold_num}_confusion_matrix.pdf'), bbox_inches='tight')
    plt.close()
    
    # 保存混淆矩阵数据
    cm_data = {
        'confusion_matrix': cm.tolist(),
        'threshold': 0.5,
        'class_names': ['Negative', 'Positive']
    }
    with open(os.path.join(plots_dir, f'fold_{fold_num}_confusion_matrix.json'), 'w') as f:
        json.dump(cm_data, f, indent=4)
    
    print(f"Fold {fold_num} 评估图表已保存至: {plots_dir}")
    
    return {
        'auc': roc_auc,
        'average_precision': average_precision,
        'confusion_matrix': cm.tolist()
    }

if __name__ == "__main__":
    set_seed(42)

    # 定义所有PDB ID
    all_pdbs = [
        '1A0S', '1A12', '1A3F', '1AFB', '1AHS', '1B08', '1BAW', '1C28', '1C3Q', '1C3X',
        '1C5E', '1CE0', '1D5F', '1DBF', '1DD1', '1EF8', '1EL6', '1F7Q', '1FTH', '1FXZ',
        '1GCM', '1HX6', '1IQA', '1J5S', '1JLJ', '1JS0', '1KI9', '1KRV', '1MLV', '1MPR',
        '1O7K', '1O8O', '1O91', '1ODE', '1OTG', '1PB0', '1PPR', '1QBZ', '1RGX', '1SED',
        '1SG2', '1SG4', '1SGJ', '1SJN', '1TD3', '1TGG', '1U5Y', '1UFL', '1V4N', '1V6H',
        '1VL0', '1VMF', '1VMK', '1W9Z', '1WBH', '1WRV', '1WT6', '1WU8', '1XBF', '1XHO',
        '1XRG', '1Y4M', '1Y8T', '1YB0', '1YQ6', '1YQQ', '1YU4', '1YX1', '1ZVB', '2A5Z',
        '2AKF', '2B4I', '2B9B', '2BHW', '2BSD', '2CHC', '2CU5', '2CW5', '2CZ4', '2D39',
        '2DQY', '2E2A', '2E66', '2EBO', '2EKM', '2EKN', '2F0C', '2F7Y', '2FB5', '2FE8',
        '2FVH', '2GTR', '2GUM', '2GVH', '2H6L', '2I9D', '2IEQ', '2INU', '2IUA', '2IUM',
        '2J9E', '2JCA', '2NWL', '2O4V', '2O66', '2O8X', '2OBE', '2OKD', '2OTM', '2P2L',
        '2P90', '2PBZ', '2Q01', '2Q0T', '2RE9', '2RGQ', '2TNF', '2UZH', '2V2H', '2VES',
        '2VRS', '2WLG', '2WPS', '2WPZ', '2WQ4', '2WRC', '2WW2', '2WW6', '2X3H', '2XC1',
        '2XGF', '2XU8', '2Y8C', '2YGC', '2YKO', '2YNY', '2YNZ', '2YZJ', '2ZBV', '2ZFC',
        '3A76', '3B7K', '3BHP', '3BSY', '3C6V', '3CJ8', '3CM1', '3CPX', '3CQO', '3CYN',
        '3D9X', '3DA0', '3DC7', '3DLI', '3EH0', '3EMF', '3EXW', '3F4F', '3FLH', '3FOB',
        '3G64', '3GDC', '3GKB', '3GOS', '3GQH', '3GTZ', '3H35', '3H6X', '3H81', '3HEZ',
        '3HFE', '3HTN', '3HYK', '3I3F', '3IRS', '3IWT', '3JQY', '3JS9', '3K0T', '3LAO',
        '3LGI', '3LKJ', '3MAE', '3MBQ', '3MOG', '3N27', '3N4G', '3NCR', '3NE1', '3NSG',
        '3NTN', '3O0Y', '3OL0', '3OPK', '3QK3', '3QLL', '3QXZ', '3R1W', '3R5C', '3R9Q',
        '3RMR', '3RWN', '3SOZ', '3SWF', '3SWY', '3TAS', '3TIS', '3TY1', '3UIA', '3VBM',
        '3VNP', '3W93', '3WJ7', '3WPR', '3WQA', '3X2Z', '4A0U', '4ADG', '4AE2', '4B6R',
        '4C8H', '4C8S', '4CE7', '4CGC', '4CO4', '4CQ6', '4D3H', '4DGQ', '4DI1', '4DYS',
        '4DZN', '4E38', '4E98', '4E9W', '4EB8', '4EDI', '4EZD', '4FAY', '4FIO', '4G1A',
        '4G2K', '4GJH', '4GVS', '4H41', '4HI1', '4HUS', '4I6V', '4IC5', '4IYQ', '4JCU',
        '4JDN', '4JJ2', '4JJ9', '4JQS', '4K29', '4K6U', '4KG8', '4KNT', '4LEH', '4LGO',
        '4LK5', '4MEJ', '4MI2', '4MOU', '4N23', '4N72', '4NBQ', '4NCV', '4NNQ', '4OJL',
        '4OOP', '4OXM', '4OZN', '4Q1I', '4R7T', '4RFU', '4RHM', '4RU5', '4RWX', '4RX6',
        '4U18', '4U5R', '4UOF', '4USH', '4USX', '4UW7', '4UXE', '4UXX', '4WYJ', '4X3N',
        '4XC5', '4XL8', '4YOU', '4YV3', '5APQ', '5APY', '5CUM', '5DRK', '5DUC', '5EIL',
        '5ERO', '5EUR', '5F1C', '5FUS', '5G47', '5G5N', '5GYL', '5H1X', '5HBA', '5HP8',
        '5I2M', '5JBX', '5JQM', '5JS4', '5K21', '5KC6', '5LDT', '5M9F', '5N83', '5NS8',
        '5NSW', '5NXR', '5O34', '5O65', '5OMI', '5STD', '5T9Y', '5TB7', '5TDG', '5TOH',
        '5TOI', '5U5B', '5UCQ', '5UN0', '5UXT', '5V0Z', '5V13', '5VMK', '5VXB', '5VXO',
        '5W0B', '5W6H', '5XL8', '5XU7', '5Y5Q', '5YBZ', '5Z1Q', '5Z81', '5ZUV', '6AHG',
        '6ARB', '6C72', '6CUQ', '6DK0', '6DZ3', '6EGN', '6EKE', '6EU4', '6EUA', '6EUN',
        '6G43', '6GAJ', '6GAP', '6GDV', '6H21', '6H9M', '6HBE', '6HDE', '6JYY', '6L6N',
        '6L8P', '6LJ3', '6MGE', '6MSR', '6NTB', '6NW9', '6NYI', '6OSS', '6OZB', '6PUA',
        '6QP4', '6R5W', '6RX3', '6RZ2', '6S8C', '6SSW', '6TGF', '6TXD', '6U66', '6V78',
        '6VIS', '6VVR', '6WE5', '6WH5', '6WV2', '6X7Q', '6YC6', '6YW9', '6ZZL', '7BRU',
        '7BW9', '7DOV', '7DSZ', '7FE0', '7JM1', '7L33', '7O92', '7O9V', '7OAH', '7OJ6',
        '7P3H', '7P4L', '7QDK', '7QWC', '7R1M', '7R5Z', '7VVW', '7W1F', '7WLS', '7ZEA',
        '7ZSQ', '8ABU', '8AHZ', '8AP3', '8DQ6', '8F3B', '8GPP', '8K06', '8K2Y', '8OFS',
        '8OKS', '8OKW', '8OML', '8ONF', '8Q4E', '8QEU', '8S38', '8SJI', '8SWU', '8T9Q',
        '8U2V', '8UUS', '8UZ4', '8VJ2', '8WCO', '8XHU', '8YKC', '9BKB', '9FAG',
        '1C4T', '1CBU', '1JXZ', '1K4M', '1KKE', '1LUA', '1MWW', '1MZZ', '1N3B', '1QDN',
        '1THJ', '1Y5E', '2IEX', '2IG8', '2Q6E', '2W95', '2WAM', '3D5K', '3EMO', '3FUY',
        '3HYT', '3IAC', '3MCH', '3PR7', '3VDL', '3ZJB', '4AOZ', '4C46', '4CQJ', '4DME',
        '4F2D', '4M99', '4NF2', '4UE0', '4WOL', '5B8F', '5E1T', '5J0J', '5KA6', '5ONU',
        '5SUV', '6C7C', '6PNZ', '7BBZ', '7EEA', '7M58', '8AU0', '8HSN', '8JUR', '8QKY'
    ]
    
    dataset = ResidueDataset(
        esm_feature_dir="/fs/fast/u2023000745/work_2/homo_c3_trimer/1_6_prott5/prott5_same_esm2_449",
        label_csv="/fs/fast/u2023000745/work_2/homo_c3_trimer/output_01_from_fasta.csv",
        pdb_id_list=all_pdbs,
        window_size=2
    )
    
    # 使用分层抽样方法获取折索引
    fold_indices = get_stratified_fold_indices(dataset)
    
    all_metrics = {
        'accuracy': [],
        'f1': [],
        'auc': [],
        'top1': [],
        'top5': [],
        'top10': [],
        'top20': [],
        'top50': [],
        'top100': [],
    }
    
    # 创建总体结果文件
    overall_metrics_file = "cross_val_results_mlp_res_optuna.txt"
    with open(overall_metrics_file, "w") as f:
        f.write("Cross-Validation Results (MLP+RES with Optuna)\n")
        f.write("="*50 + "\n")
        f.write("Fold  Accuracy  F1-Score  AUC      Top1%    Top5%    Top10%   Top20%   Top50%   Top100%\n")
        f.write("="*50 + "\n")
    
    # 存储最佳超参数
    best_params_all_folds = {}
    
    for fold in range(NUM_FOLDS):
        print(f"\n\n======= STARTING FOLD {fold+1}/{NUM_FOLDS} =======")
        
        # Split data for this fold (保持不变)
        test_pdbs = fold_indices[fold]
        train_val_pdbs = [pdb for i, pdb_list in enumerate(fold_indices) for pdb in pdb_list if i != fold]
        
        # Split train_val_pdbs into train and validation
        train_pdbs, val_pdbs = train_test_split(
            train_val_pdbs, test_size=0.1, random_state=42
        )
        
        train_indices = split_dataset_by_pdb(dataset, train_pdbs)
        val_indices = split_dataset_by_pdb(dataset, val_pdbs)
        test_indices = split_dataset_by_pdb(dataset, test_pdbs)
        
        # ==== 根据开关决定是否进行下采样 ====
        if ENABLE_DOWNSAMPLING:
            # 获取训练集的标签
            train_labels = [dataset.samples[i]['label'] for i in train_indices]
            
            # 统计正负样本数量
            positive_count = sum(train_labels)
            negative_count = len(train_labels) - positive_count
            
            # 计算需要保留的负样本数量（目标比例 4:100）
            target_negative_count = int(positive_count * 100 / 18)
            
            # 创建索引数组
            all_train_indices = np.array(train_indices)
            positive_indices = all_train_indices[np.where(np.array(train_labels) == 1)[0]]
            negative_indices = all_train_indices[np.where(np.array(train_labels) == 0)[0]]
            
            # 随机选择负样本以满足比例要求
            if len(negative_indices) > target_negative_count:
                np.random.shuffle(negative_indices)
                selected_negative_indices = negative_indices[:target_negative_count].tolist()
            else:
                selected_negative_indices = negative_indices.tolist()
            
            # 合并正样本和选择的负样本
            downsampled_train_indices = positive_indices.tolist() + selected_negative_indices
            
            # 打印采样信息
            print(f"Fold {fold+1} - Downsampling enabled")
            print(f"Fold {fold+1} - Before downsampling: Pos: {positive_count}, Neg: {negative_count}")
            print(f"Fold {fold+1} - After downsampling: Pos: {len(positive_indices)}, Neg: {len(selected_negative_indices)}")
            print(f"Fold {fold+1} - New sample ratio: 1:{len(selected_negative_indices)/len(positive_indices):.2f}")
            
            train_indices_to_use = downsampled_train_indices
        else:
            print(f"Fold {fold+1} - Downsampling disabled")
            train_indices_to_use = train_indices
        
        # 创建数据集子集
        train_dataset = torch.utils.data.Subset(dataset, train_indices_to_use)
        val_dataset = torch.utils.data.Subset(dataset, val_indices)
        test_dataset = torch.utils.data.Subset(dataset, test_indices)
        
        # 创建DataLoader
        train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, pin_memory=True)
        val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE*2, pin_memory=True)
        
        # ==== Optuna超参数优化 ====
        print(f"开始Optuna超参数优化 (Fold {fold+1})...")
        study = optuna.create_study(
            direction='maximize',
            sampler=TPESampler(seed=42)
        )
        
        # 使用lambda函数传递额外的参数
        study.optimize(
            lambda trial: objective(trial, train_loader, val_loader),
            n_trials=OPTUNA_N_TRIALS,
            show_progress_bar=True
        )
        
        # 获取最佳参数
        best_params = study.best_params
        best_value = study.best_value
        
        best_params_all_folds[f'fold_{fold+1}'] = best_params
        
        print(f"Fold {fold+1} 最佳参数: {best_params}")
        print(f"Fold {fold+1} 最佳AUC: {best_value:.6f}")

        # 使用最佳参数训练最终模型之前，构建hidden_dims参数
        if 'hidden_dim_base' in best_params:
            # 构建hidden_dims列表
            hidden_dims = []
            current_dim = best_params['hidden_dim_base']
            for i in range(best_params['hidden_layers']):
                hidden_dims.append(int(current_dim))
                current_dim = current_dim * best_params['reduction_factor']
            best_params['hidden_dims'] = hidden_dims

        model = ContactPredictor(best_params)
        print(f"使用最佳参数训练最终模型 (Fold {fold+1})...")
        model.fit(train_loader, val_loader) 

        
        # 获取测试数据和预测结果
        X_test = []
        y_test = []
        for data in test_dataset:
            X_test.append(data[0].numpy())
            y_test.append(data[1].item())
        X_test = np.array(X_test)
        y_test = np.array(y_test)
        
        test_preds = model.predict(X_test)
        test_probs = model.predict_proba(X_test)
        
        print(f"测试特征维度: {X_test.shape}") 
        
        # 计算指标
        test_metrics = {
            'accuracy': accuracy_score(y_test, test_preds),
            'f1': f1_score(y_test, test_preds),
            'auc': roc_auc_score(y_test, test_probs),
        }
        
        top1_acc = calculate_top1_accuracy(dataset, test_indices, test_probs)
        top5_acc = calculate_top5_accuracy(dataset, test_indices, test_probs)
        top10_acc = calculate_top10_accuracy(dataset, test_indices, test_probs)
        top20_acc = calculate_top20_accuracy(dataset, test_indices, test_probs)
        top50_acc = calculate_top50_accuracy(dataset, test_indices, test_probs)
        top100_acc = calculate_top100_accuracy(dataset, test_indices, test_probs)
        
        # 保存折叠指标
        all_metrics['accuracy'].append(test_metrics['accuracy'])
        all_metrics['f1'].append(test_metrics['f1'])
        all_metrics['auc'].append(test_metrics['auc'])
        all_metrics['top1'].append(top1_acc)
        all_metrics['top5'].append(top5_acc)
        all_metrics['top10'].append(top10_acc)
        all_metrics['top20'].append(top20_acc)
        all_metrics['top50'].append(top50_acc)
        all_metrics['top100'].append(top100_acc)
        
        # 创建折叠输出目录
        fold_output_dir = f"fold_{fold+1}_mlp_res_optuna"
        os.makedirs(fold_output_dir, exist_ok=True)
        
        # 绘制和保存评估图表
        print(f"绘制Fold {fold+1}评估图表...")
        plot_metrics = plot_and_save_evaluation_plots(y_test, test_probs, fold_output_dir, fold+1)
        
        # 保存模型
        torch.save({
            'model_state_dict': model.model.state_dict(),
            'params': best_params
        }, f"{fold_output_dir}/mlp_res_model.pth")

        # 保存Optuna研究
        with open(f"{fold_output_dir}/optuna_study.pkl", "wb") as f:
            pickle.dump(study, f)
        
        # 保存测试结果
        evaluate_with_pdb_output(dataset, test_indices, test_probs, output_dir=os.path.join(fold_output_dir, "test_results"))
        
        # 保存指标到JSON文件
        fold_metrics = {
            'test_probs': test_probs.tolist(),
            'test_indices': test_indices,
            'test_labels': y_test.tolist(),
            'top1_accuracy': top1_acc,
            'top5_accuracy': top5_acc,
            'top10_accuracy': top10_acc,
            'top20_accuracy': top20_acc,
            'top50_accuracy': top50_acc,
            'top100_accuracy': top100_acc,
            'auc': test_metrics['auc'],
            'f1': test_metrics['f1'],
            'best_params': best_params,
            'best_auc': best_value,
            'plot_metrics': plot_metrics
        }
        
        with open(f"{fold_output_dir}/fold_results.json", "w") as f:
            json.dump(fold_metrics, f, indent=4)
            
        # 保存数据拆分信息
        split_info = {
            "train_indices": train_indices_to_use,
            "val_indices": val_indices,
            "test_indices": test_indices,
            "train_pdbs": train_pdbs,
            "val_pdbs": val_pdbs,
            "test_pdbs": test_pdbs,
            "downsampling_enabled": ENABLE_DOWNSAMPLING
        }
        
        with open(f"{fold_output_dir}/data_split.pkl", "wb") as f:
            pickle.dump(split_info, f)
            
        # 保存指标明细到TXT文件
        fold_metrics_txt = os.path.join(fold_output_dir, "fold_metrics.txt")
        with open(fold_metrics_txt, "w") as f:
            f.write(f"Fold {fold+1} Results (MLP+RES with Optuna)\n")
            f.write("="*40 + "\n")
            f.write(f"Downsampling enabled: {ENABLE_DOWNSAMPLING}\n")
            f.write(f"Optuna最佳AUC: {best_value:.6f}\n")
            f.write(f"Accuracy:    {test_metrics['accuracy']:.6f}\n")
            f.write(f"F1-Score:    {test_metrics['f1']:.6f}\n")
            f.write(f"AUC:         {test_metrics['auc']:.6f}\n")
            f.write(f"Top1 Acc:    {top1_acc:.6f}\n")
            f.write(f"Top5 Acc:    {top5_acc:.6f}\n")
            f.write(f"Top10 Acc:   {top10_acc:.6f}\n")
            f.write(f"Top20 Acc:   {top20_acc:.6f}\n")
            f.write(f"Top50 Acc:   {top50_acc:.6f}\n")
            f.write(f"Top100 Acc:  {top100_acc:.6f}\n")
            f.write("="*40 + "\n")
            f.write(f"Number of test samples: {len(test_indices)}\n")
            f.write(f"Number of test PDBs: {len(set([dataset.samples[i]['pdb_id'] for i in test_indices]))}\n")
            f.write(f"Positive samples in test set: {sum(y_test)}\n")
            f.write(f"Negative samples in test set: {len(y_test) - sum(y_test)}\n")
            
            # 添加训练信息摘要
            f.write("\nTraining Summary:\n")
            f.write(f"Number of training samples: {len(train_indices_to_use)}\n")
            f.write(f"Number of training PDBs: {len(train_pdbs)}\n")
            f.write(f"Number of validation samples: {len(val_indices)}\n")
            f.write(f"Number of validation PDBs: {len(val_pdbs)}\n")
            f.write(f"Optuna trials: {OPTUNA_N_TRIALS}\n")
        
        # 将指标添加到总体结果文件中
        with open(overall_metrics_file, "a") as f:
            f.write(f"{fold+1:4d}  {test_metrics['accuracy']:.6f}  {test_metrics['f1']:.6f}  " 
                    f"{test_metrics['auc']:.6f}  {top1_acc:.6f}  {top5_acc:.6f}  "
                    f"{top10_acc:.6f}  {top20_acc:.6f}  {top50_acc:.6f}  {top100_acc:.6f}\n")
        
        print(f"\nCompleted fold {fold+1}/{NUM_FOLDS}")
    
    # 计算和保存总体统计结果
    with open(overall_metrics_file, "a") as f:
        f.write("\nSummary Statistics\n")
        f.write("="*50 + "\n")
        f.write("Metric      Mean      Std\n")
        f.write("="*50 + "\n")
        
        for metric in all_metrics:
            values = all_metrics[metric]
            mean = np.mean(values)
            std = np.std(values)
            f.write(f"{metric:<8}  {mean:.6f}  {std:.6f}\n")
    
    # 保存所有指标到JSON文件
    with open("cross_val_results_mlp_res_optuna.json", "w") as f:
        json.dump(all_metrics, f, indent=4)
    print("\nCross-validation results saved to cross_val_results_mlp_res_optuna.json")
    
    # 保存所有fold的最佳参数
    with open("best_params_all_folds.json", "w") as f:
        json.dump(best_params_all_folds, f, indent=4)
    print("Best parameters for all folds saved to best_params_all_folds.json")
    
    # 保存实验配置
    experiment_config = {
        "enable_downsampling": ENABLE_DOWNSAMPLING,
        "stratified_bins": STRATIFIED_BINS,
        "batch_size": BATCH_SIZE,
        "epochs": EPOCHS,
        "learning_rate": LEARNING_RATE,
        "num_folds": NUM_FOLDS,
        "model_type": "MLP+RES",
        "optuna_trials": OPTUNA_N_TRIALS
    }
    with open("experiment_config_mlp_res_optuna.json", "w") as f:
        json.dump(experiment_config, f, indent=4)
    print("Experiment configuration saved to experiment_config_mlp_res_optuna.json")
    
    # 计算总运行时间
    end_time = time.time()
    total_time = end_time - start_time
    hours = int(total_time // 3600)
    minutes = int((total_time % 3600) // 60)
    seconds = total_time % 60
    
    print(f"\nTotal running time: {hours}h {minutes}m {seconds:.2f}s")
    print("10-fold cross validation with Optuna optimization completed successfully!")