main_abulation.py

import csv
import gc
import logging
import math
import os, sys
import random
import time
from tqdm import tqdm
from tqdm.contrib.logging import logging_redirect_tqdm
from functools import lru_cache
from collections import deque
from typing import Literal
import torch.nn as nn


import optuna
import pandas as pd
import dgl
import numpy as np
import torch
from torch.utils.data import DataLoader, TensorDataset, Subset
import torch.nn.functional as F
import torch.optim as optim
import torch.optim.lr_scheduler as lr_scheduler
from torch.cuda.amp import autocast, GradScaler
from torch_scatter import scatter_mean, scatter_add
from pytorch_metric_learning import losses
from src.utils import set_logging, to_numpy
from src.misc.revgat.loss import loss_kd_only
from src.model.lm_gnn import RevGAT, HSAGE, MLP, global_entropy, graph_reconstruction_loss
from src.model import get_model_class
from src.args import parse_args, save_args
import src.lora as lora
from src.trainer import get_trainer_class, TextDataset, Code_Trainer
from sklearn.metrics import accuracy_score, f1_score, precision_score, recall_score, roc_auc_score, matthews_corrcoef,jaccard_score
from data_loader import CodeEvaluator,GraphTextDataset,collate_fn
from torch_scatter import scatter_add
from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops
from src.utils import search_best_threshold, search_best_alpha_threshold


logger = logging.getLogger(__name__)
logging.getLogger("torch").setLevel(logging.ERROR)
# os.environ["DGL_WORKER_CPU_PINNING"] = "0"
def worker_init_fn(worker_id):
    sys.stderr = open(os.devnull, 'w')  # 屏蔽 worker 进程的 stderr
    sys.stdout = open(os.devnull, 'w')  # 屏蔽 stdout
    np.random.seed(42 + worker_id)
    random.seed(42 + worker_id)
    torch.manual_seed(42 + worker_id)

def seed(seed=0):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    dgl.random.seed(seed)

def check_nan(model):
    for name, param in model.named_parameters():
        if param.requires_grad and torch.isnan(param).any():
            raise ValueError(f"NaN detected in parameter {name}")
        if param.grad is not None and torch.isnan(param.grad).any():
            raise ValueError(f"NaN detected in gradient of {name}")
        
def top2n_recall(y_true, y_prob, k=2):
    n_pos = sum(y_true)
    if n_pos > 0:
        top_k = math.ceil(k * n_pos)
        topk_indices = np.argsort(y_prob)[-top_k:]
        true_positives = sum(y_true[i] for i in topk_indices)
        recall = true_positives / n_pos if n_pos > 0 else 0
    else:
        recall = 0
    return recall
        
def topkk_recall(y_true, y_prob, k=5):
    n_pos = sum(y_true)
    if n_pos > 0:
        top_k = n_pos+k
        topk_indices = np.argsort(y_prob)[-top_k:]
        true_positives = sum(y_true[i] for i in topk_indices)
        recall = true_positives / n_pos 
    else:
        recall = 0
    return recall

def collate_fn1(batch):
    # 初始化空列表，用于存储批次数据
    graphs = []
    input_ids_list = []
    attention_mask_list = []
    labels = []

    # 遍历批次中的每个样本
    for item in batch:
        graphs.append(item['graph'])
        input_ids_list.append(item['input_ids'])
        attention_mask_list.append(item['attention_mask'])
        labels.append(item['labels'])  # 每个样本的label是一个张量

    # 拼接图数据
    batched_graph = dgl.batch(graphs)

    # 拼接文本数据
    input_ids = torch.cat(input_ids_list, dim=0)  # 纵向拼接input_ids
    attention_mask = torch.cat(attention_mask_list, dim=0)  # 纵向拼接attention_mask

    # 拼接labels
    labels = torch.cat(labels, dim=0)  # 将labels拼接成一个张量

    # 返回批次数据
    return {
        'graph': batched_graph,
        'input_ids': input_ids,
        'attention_mask': attention_mask,
        'labels': labels
    }

def multi_hop_self_agg(grad, edge_index, num_hops=1, alpha=0.5):
    """
    使用自身 + 多跳邻居梯度估计进行梯度补偿

    参数:
    - grad: (num_nodes, feature_dim) 原始梯度张量
    - edge_index: (2, num_edges) 图的边索引
    - num_hops: 传播的跳数
    - alpha: 邻居梯度衰减因子（仅适用于 num_hops > 1）

    返回:
    - compensated_grad: (num_nodes, feature_dim) 经过补偿后的梯度
    """
    num_nodes = grad.shape[0]
    compensated_grad = grad.clone()  # 先初始化为自身梯度
    node_weight = torch.ones((num_nodes, 1), device=grad.device)  # 记录权重归一化因子
    current_grad = grad.clone()  # 作为初始梯度
    
    for hop in range(num_hops):
        # 计算一跳邻居梯度均值
        neighbor_grad = torch.zeros_like(grad)
        neighbor_weight = torch.zeros((num_nodes, 1), device=grad.device)
        
        # 聚合邻居的梯度
        scatter_add(current_grad[edge_index[0]], edge_index[1], dim=0, out=neighbor_grad)
        scatter_add(torch.ones_like(node_weight)[edge_index[0]], edge_index[1], dim=0, out=neighbor_weight)
        
        # 归一化：如果节点没有邻居，避免除以 0
        neighbor_weight = torch.where(neighbor_weight > 0, neighbor_weight, torch.ones_like(neighbor_weight))
        neighbor_grad /= neighbor_weight  # 计算均值
        
        # 计算当前 hop 的补偿
        compensated_grad += alpha**hop * neighbor_grad
        node_weight += alpha**hop  # 累计归一化因子
        
        # 更新 current_grad 为本次邻居梯度，以便用于下一跳
        current_grad = neighbor_grad.clone()

    # 最终归一化
    compensated_grad /= node_weight
    return compensated_grad

def multi_hop_aggregation(grad, edge_index, num_nodes, num_hops=1, decay=0.5):
    """
    对所有节点进行多跳梯度聚合。
    当 num_hops=1 时，直接聚合一跳邻居的梯度信息。
    
    参数:
        grad: 原始梯度信息，形状为 [num_nodes, feature_dim]
        edge_index: 图的边索引，形状为 [2, E]，第一行为源节点，第二行为目标节点
        num_nodes: 节点总数
        num_hops: 聚合跳数，>=1 表示至少聚合一跳邻居信息
        decay: 衰减因子，用于降低多跳传播时越远节点的影响
    返回:
        聚合后的梯度，形状为 [num_nodes, feature_dim]
    """
    aggregated_grad = 0  # 用于累加每一跳的聚合结果
    current_grad = grad  # 初始梯度

    for hop in range(num_hops):
        # 每一次循环将 current_grad 聚合到目标节点上
        current_grad = scatter_mean(current_grad[edge_index[0]], edge_index[1], dim=0, dim_size=num_nodes)
        aggregated_grad += (decay ** hop) * current_grad
    return aggregated_grad


class FocalLoss(torch.nn.Module):
    def __init__(self,
            alpha=None,
            class_counts=None,
            alpha_smoothing='log',  # 'none', 'log', 'sqrt'
            normalize_alpha=True,   # 是否归一化为均值 1
            gamma=1.0,
            smoothing=0.1,
            reduction='mean'):
        super().__init__()
        self.gamma = gamma
        self.smoothing = smoothing
        self.reduction = reduction

        if alpha is None and class_counts is not None:
            counts = torch.tensor(class_counts, dtype=torch.float32)
            total = counts.sum()
            alpha = total / counts

            if alpha_smoothing == 'log':
                alpha = torch.log1p(alpha)
                alpha = torch.sqrt(alpha)
            elif alpha_smoothing == 'sqrt':
                alpha = torch.sqrt(alpha)

            if normalize_alpha:
                alpha = alpha / alpha.mean()

        if alpha is not None:
            self.register_buffer('alpha', alpha)
        else:
            self.alpha = None

    def forward(self, inputs, targets):
        os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
        num_classes = inputs.size(1)
        device = inputs.device

        with torch.no_grad():
            true_dist = torch.full_like(inputs, self.smoothing / (num_classes - 1))
            true_dist.scatter_(1, targets.unsqueeze(1), 1.0 - self.smoothing)

        log_probs = F.log_softmax(inputs, dim=1)
        probs = log_probs.exp()
        pt = torch.sum(probs * true_dist, dim=1)
        log_pt = torch.sum(log_probs * true_dist, dim=1)

        if self.alpha is not None:
            assert targets.dtype == torch.long, f"targets.dtype={targets.dtype}, must be torch.long"
            assert targets.min() >= 0, f"target min = {targets.min().item()} < 0"
            assert targets.max() < num_classes, f"target max = {targets.max().item()} >= num_classes={num_classes}"
            alpha_t = self.alpha.to(device)[targets]
        else:
            alpha_t = 1.0

        loss = -alpha_t * (1 - pt) ** self.gamma * log_pt
        
        os.environ["CUDA_LAUNCH_BLOCKING"] = "0"
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        else:
            return loss
        


class ReplaceRowsFunction(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, rows, replacement, num_nodes, edge_index):
        """
          input: 原始特征矩阵，shape为 [N, F]
          rows: 被替换行的索引
          replacement: 替换的特征，这部分节点需要梯度
          edge_index: 图的边索引，shape为 [2, E]
          num_nodes: 节点总数
        """
        ctx.save_for_backward(input, rows, replacement)
        ctx.edge_index = edge_index
        ctx.num_nodes = num_nodes
        output = input.clone()
        output[rows] = replacement
        return output

    @staticmethod
    def backward(ctx, grad_output):
        input, rows, replacement = ctx.saved_tensors
        edge_index = ctx.edge_index
        num_nodes = ctx.num_nodes
        grad_input = grad_output.clone()
        grad_input[rows] = 0
        alpha = 0.05  # 可调超参数，调节补偿力度
        
        # # 全图平均估计  采样消融可以使用
        # mask = torch.ones(grad_output.shape[0], dtype=torch.bool, device=grad_output.device)
        # mask[rows] = False
        # # 计算未替换节点的全局平均梯度
        # if mask.sum() > 0:
        #     global_avg = grad_output[mask].mean(dim=0)
        # else:
        #     global_avg = torch.zeros_like(grad_output[0])
        # grad_replacement = grad_output[rows] + alpha * global_avg
        
        # grad_replacement = grad_output[rows].clone()

        # 局部邻居估计
        grad_neighbor = multi_hop_aggregation(grad_output, edge_index, num_nodes, 2, alpha)
        grad_replacement = 1/(1+alpha)*(grad_output[rows] + alpha * grad_neighbor[rows])

        #grad_replacement = grad_output[rows]

        return grad_input, None, grad_replacement, None, None
    
    # #可以考虑补偿-qianwen 函数
    # @staticmethod
    # def backward(ctx, grad_output):
    #     input, rows, replacement = ctx.saved_tensors
    #     edge_index = ctx.edge_index
    #     num_nodes = ctx.num_nodes

    #     # === 明确两个超参数 ===
    #     decay = 0.5   # 多跳内部衰减（控制 G1, G2, ... 的权重）
    #     beta = 0.1    # 最终补偿强度（控制原始梯度和平滑梯度的混合比例）

    #     grad_input = grad_output.clone()
    #     grad_input[rows] = 0

    #     # 调用时显式指定参数名，避免混淆
    #     grad_neighbor = multi_hop_aggregation(
    #         grad=grad_output,
    #         edge_index=edge_index,
    #         num_nodes=num_nodes,
    #         num_hops=2,
    #         decay=decay  # ← 不再用 alpha！
    #     )

    #     # 使用清晰的凸组合
    #     grad_replacement = (1 - beta) * grad_output[rows] + beta * grad_neighbor[rows]

    #     return grad_input, None, grad_replacement, None, None
    
replace_rows = ReplaceRowsFunction.apply

class BiDirectNeighborSampler(dgl.dataloading.NeighborSampler):
    def sample_neighbors(self, g, nodes):
        in_neighbors = g.in_edges(nodes)[0]  # 采样入度邻居
        out_neighbors = g.out_edges(nodes)[1]  # 采样出度邻居
        all_neighbors = torch.cat([in_neighbors, out_neighbors])  # 合并
        return g.subgraph(all_neighbors)

class MeanAggregation(MessagePassing):
    def __init__(self):
        super(MeanAggregation, self).__init__(aggr='mean')

    def forward(self, x, edge_index):
        # x has shape [N, in_channels]
        # edge_index has shape [2, E]

        # Step 1: Add self-loops to the adjacency matrix.
        _edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))

        # Step 4-5: Start propagating messages.
        return self.propagate(_edge_index, x=x)

def compute_jsd(dist1, dist2):
    dist_mean = (dist1 + dist2) / 2.
    jsd = (F.kl_div(dist_mean.log(), dist1, reduction='none') + 
           F.kl_div(dist_mean.log(), dist2, reduction='none')) / 2.
    return jsd

class LM_GNN():
    def __init__(self, args, **kwargs) -> None:
        # 参数初始化
        self.args = args
        self.epsilon = args.eps if args.eps is not None else 1 - math.log(2)
        self.device = torch.device(f"cuda:{args.gpu}" if not args.cpu else "cpu")  # 设备设置
        self.dtype = torch.float16 if args.fp16 else torch.float32  # 数据类型选择
        
        # 数据相关
        self.data=[] #存储富文本图数据
        self.graphs = []  # 存储所有的图
        self.text_tokens = []  # 存储所有图的tokenized文本数据
        self.feat_static = None  # 图的静态特征
        # self.feat = None
        self.labels = None  # 存储所有图的标签数据
        self.uids = None
        self.train_idx = None  # 存储所有训练集索引
        self.val_idx = None  # 存储所有验证集索引
        self.test_idx = None  # 存储所有测试集索引
        self.split_idx = None  # 数据集划分信息（train, valid, test）
        self.train_dataset = None
        self.test_dataset = None
        #self.preds = None

        # 模型相关
        self.model_lm = None  # 语言模型
        self.model_gnn = None  # 图神经网络模型
        self.is_lm = []  # 语言模型参数标记
        self.require_grad = []  # 需要梯度更新的参数
        self.lora_added = False  # 是否使用LoRA

        # 训练相关
        self.optimizer = None  # 优化器
        self.criterion = None  # 损失函数
        self.NTXentLoss = losses.NTXentLoss(temperature = self.args.label_smoothing)
        self.SupConLoss = losses.SupConLoss(temperature = self.args.label_smoothing)
        self.evaluator = None  # 模型评估器，仅lm预训练使用
        self.whole_graph = False  # 是否使用整个图

        # 试验相关
        self.trial = kwargs.pop("trial", None)  # 试验编号
        
        # 辅助信息
        self.n_node = 0  # 节点数（可以根据需要更新）
        self.gpt_preds = None  # GPT模型预测（如果有）

        # GradScaler（用于FP16训练）
        self.scaler = GradScaler() if self.args.fp16 else None

        # 设备转移标记
        self.args.grad_padding2 = 0

        self.train_data_loader = None

        if(self.args.lm_x):
            self.node_classifier = nn.Linear(self.args.adapter_hidden_size, 2).to(self.device)
        
    def preprocess(self):
        # global n_node_feats

        # make bidirected
        # feat = graph.ndata["feat"]
        for i, g in enumerate(self.graphs):
            
            g = g.remove_self_loop()
            ne = g.num_edges()
            g = dgl.add_reverse_edges(g, copy_ndata=True, copy_edata=True)
            g = g.add_self_loop()
            etype = torch.zeros(g.num_edges(), dtype=torch.long)
            etype[ne:ne*2] = 1
            etype[ne*2:] = 2
            g.edata["type"] = etype
            self.graphs[i] = g 
            # logger.info(f"Total edges after adding self-loop {self.graph.number_of_edges()}")

        # self.graph.create_formats_()
        self.args.n_edge_types = 3

    def prepare(self):
        '''device, scaler, criterion, metrics'''
        if self.args.cpu:
            self.device = torch.device("cpu")
        else:
            self.device = torch.device(f"cuda:{self.args.gpu}")
        # self.labels, self.val_idx, self.test_idx = map(
        # lambda x: x.to(self.device), (self.labels, self.val_idx, self.test_idx)
        # )
        self.labels = self.labels.to(self.device)
        # 初始化GradScaler
        self.scaler = GradScaler() if self.args.fp16 else None
        
        num_total = self.data['attr']['n_nodes']
        num_pos = self.data['attr']['n_vul_nodes']
        # self.args.loss_weight = torch.tensor([weight_neg, weight_pos], dtype=torch.float32)
        # self.criterion = torch.nn.CrossEntropyLoss(label_smoothing=self.args.label_smoothing, reduction =self.args.loss_reduction)
        # self.criterion = torch.nn.CrossEntropyLoss(label_smoothing=self.args.label_smoothing, reduction =self.args.loss_reduction, weight=self.args.loss_weight)
        self.criterion = FocalLoss(class_counts=[num_total-num_pos, num_pos], 
                                   alpha_smoothing='log',  
                                   gamma=self.args.focal_loss_gamma,
                                   smoothing=self.args.label_smoothing_factor, 
                                   reduction=self.args.loss_reduction,
                                   normalize_alpha = False
                                   #normalize_alpha = True
                                   )
        self.metrics = {
            'acc' : lambda y_true,y_pred,y_prob: accuracy_score(y_true, y_pred),
            'f1' : lambda y_true,y_pred,y_prob: f1_score(y_true, y_pred, average="weighted")\
                if self.args.num_labels > 2 else f1_score(y_true, y_pred),
            'prec': lambda y_true,y_pred,y_prob: precision_score(y_true, y_pred, average="weighted")\
                if self.args.num_labels > 2 else precision_score(y_true, y_pred),
            'rec' : lambda y_true,y_pred,y_prob: recall_score(y_true, y_pred, average="weighted")\
                if self.args.num_labels > 2 else recall_score(y_true, y_pred),
            'auc': lambda y_true,y_pred,y_prob: roc_auc_score(y_true, y_prob, multi_class="ovo", average="micro")\
                if self.args.num_labels > 2 else roc_auc_score(y_true, y_prob),
            'auc1': lambda y_true,y_pred,y_prob: roc_auc_score(y_true, y_prob, multi_class="ovo", average="micro")\
                if self.args.num_labels > 2 else roc_auc_score(y_true, y_prob),
            'mcc': lambda y_true,y_pred,y_prob: matthews_corrcoef(y_true, y_pred),
            # 'iou1':lambda y_true,y_pred: jaccard_score(y_true, y_pred, average="micro")\
            #     if self.args.num_labels > 2 else jaccard_score(y_true, y_pred),
            'iou':lambda y_true,y_pred,y_prob: jaccard_score(y_true, y_pred, average="micro")\
                if self.args.num_labels > 2 else jaccard_score(y_true, y_pred),
        }
        self.metrics_test = {
            'acc' : lambda y_true,y_pred,y_prob: accuracy_score(y_true, y_pred),
            'f1' : lambda y_true,y_pred,y_prob: f1_score(y_true, y_pred, average="weighted")\
                if self.args.num_labels > 2 else f1_score(y_true, y_pred),
            'prec': lambda y_true,y_pred,y_prob: precision_score(y_true, y_pred, average="weighted")\
                if self.args.num_labels > 2 else precision_score(y_true, y_pred),
            'rec' : lambda y_true,y_pred,y_prob: recall_score(y_true, y_pred, average="weighted")\
                if self.args.num_labels > 2 else recall_score(y_true, y_pred),
            'auc': lambda y_true,y_pred,y_prob: roc_auc_score(y_true, y_prob, multi_class="ovo", average="micro")\
                if self.args.num_labels > 2 else roc_auc_score(y_true, y_prob),
            'auc1': lambda y_true,y_pred,y_prob: roc_auc_score(y_true, y_pred, multi_class="ovo", average="micro")\
                if self.args.num_labels > 2 else roc_auc_score(y_true, y_pred),
            'mcc': lambda y_true,y_pred,y_prob: matthews_corrcoef(y_true, y_pred),
            # 'iou1':lambda y_true,y_pred: jaccard_score(y_true, y_pred, average="micro")\
            #     if self.args.num_labels > 2 else jaccard_score(y_true, y_pred),
            'iou':lambda y_true,y_pred,y_prob: jaccard_score(y_true, y_pred, average="micro")\
                if self.args.num_labels > 2 else jaccard_score(y_true, y_pred),
            'rec@2n':lambda y_true,y_pred,y_prob: top2n_recall(y_true,y_prob, k=2),
            'rec@1.5n':lambda y_true,y_pred,y_prob: top2n_recall(y_true,y_prob, k=1.5),
            'rec@3+':lambda y_true,y_pred,y_prob: topkk_recall(y_true,y_prob,k=3),
            'rec@5+':lambda y_true,y_pred,y_prob: topkk_recall(y_true,y_prob,k=5),
        }

    def save_stat(self, epoch, e2e, best, name):
        out_dir = f"{self.args.save}/ckpt"
        fname = os.path.join(out_dir, f"{name}_stat.pt")
        torch.save({
            'epoch': epoch,
            'gnn_dict': self.model_gnn.state_dict(),
            'lm_dict': self.model_lm.state_dict() if self.model_lm else None,
            'optm_dict': self.optimizer.state_dict(),
            'feat_static': self.feat_static,
            'full_ft': e2e,
            # 'args': self.args,
            'best': best
            # 可以添加其他你需要保存的状态
        }, fname)
        logger.info(f"Saving stat ckpt for {name} ...")
    
    def save_pred(self, pred, name):
        os.makedirs(f"{self.args.save}/cached_embs", exist_ok=True)
        torch.save(pred, f"{self.args.save}/cached_embs/logits_{name}.pt")
        torch.save(self.feat_static, f"{self.args.save}/cached_embs/x_embs_{name}.pt")
        logger.warning(f"Saving logits & x_embs to {self.args.save}/cached_embs/_{name}.pt")
    # @lru_cache(8)
    def cal_labels(self, length, labels, idx = None):
        '''label编码'''
        onehot = torch.zeros([length, self.args.num_labels], device=self.device, 
                            #  dtype=self.dtype
                            )
        if idx is not None and len(idx) > 0:
            onehot[idx, labels[idx]] = 1
        elif length == len(labels):
            onehot.scatter_(1, labels.unsqueeze(1), 1)
        return onehot
    
    def custom_loss(self, labels, x1, feat, batch_size=0, lambda_1=0.2, lambda_2=0.1, tau=0.1, device=None):
        """
        labels: (batch_size,)
        x1: (batch_size, embedding_dim)  # 假设 x1 是特征表示
        is_batch: 如果为True，则分批次计算LSelf和LWeakly；否则正常计算
        """
        if device == 'cpu':
            labels = labels.to('cpu')
            x1 = x1.to('cpu')
        
        # 1. Cross-Entropy Loss (LCE)
        y1 = self.criterion(x1, labels)
        # y1 = F.cross_entropy(x1, labels, reduction=self.args.loss_reduction, label_smoothing=self.args.label_smoothing)
        LCE = torch.log(self.epsilon + y1) - math.log(self.epsilon)
        # loss = LCE
        
        # 2. Self-Supervised Loss (LSelf) & Weakly-Supervised Loss (LWeakly)
        if torch.any(labels != 0):
            if batch_size > 0:
                # 分批次计算
                n = feat.size(0)
                LSelf = 0.0
                LWeakly = 0.0
                deno = 0
                for i in range(0, n, batch_size):  # 分成两批计算
                    end = min(i + batch_size, n)
                    x1_batch = feat[i:end]
                    labels_batch = labels[i:end]
                    if torch.any(labels_batch != 0):
                        LSelf += self.NTXentLoss(x1_batch, labels_batch)
                        LWeakly += self.SupConLoss(x1_batch, labels_batch)
                        deno += 1
                if deno > 0:        
                    LSelf /= deno  # 平均两批的结果
                    LWeakly /= deno  # 平均两批的结果
            else:
                #正常计算
                LSelf = self.NTXentLoss(feat, labels)
                LWeakly = self.SupConLoss(feat, labels)
            if LSelf.item() == 0.0 and LWeakly.item() == 0.0:
                loss = LCE
            else:
                loss = 0.6*LCE + lambda_1 * LSelf + lambda_2 * LWeakly
        else:
            loss = LCE
        return loss
    
    def init_loader(self):
        '''dataloader，节点采样在这里进行'''
        self.graph_loaders = []  # 用于保存每张图的 DataLoader
        self.graph_loaders_sec = []

        for batch in self.train_data_loader:
            batched_graph = batch['graph']#.to(device=self.device)
            
            if self.args.grad_padding > 0:
                # 控制采样：低->高，1随机一个节点，-1所有节点
                grad_block = [self.args.grad_k for _ in range(self.args.grad_padding)]
                
                if self.args.secsam_method == "nearby":
                    self.args.grad_padding2 = self.args.grad_padding + 1
                    grad_sec = [-1 for _ in range(self.args.grad_padding)]
                elif self.args.secsam_method == "morehop":
                    self.args.grad_padding2 = self.args.grad_padding + 2
                    grad_sec = [self.args.grad_k for _ in range(self.args.grad_padding2)]
                else:
                    grad_sec = []

                if self.args.frozen_padding >= 0: 
                    fz_block = [-1 for _ in range(self.args.frozen_padding)]
                    if self.args.grad_padding2 > 0:
                        sampler_sec = BiDirectNeighborSampler(fz_block + grad_sec)
                        graph_loader_sec = dgl.dataloading.DataLoader(
                            batched_graph, range(batched_graph.number_of_nodes()),sampler_sec,
                            batch_size=self.args.kernel_size,
                            shuffle=False,
                            drop_last=False,
                            num_workers=self.args.num_workers,
                            # persistent_workers=True,  # 避免 worker 反复重启
                            worker_init_fn=worker_init_fn,  # 让 worker 启动时重定向日志
                            # pin_memory=True
                            )
                    
                    sampler = BiDirectNeighborSampler(fz_block + grad_block)
                    graph_loader = dgl.dataloading.DataLoader(
                        batched_graph,range(batched_graph.number_of_nodes()),sampler,
                        batch_size=self.args.kernel_size,
                        shuffle=False,
                        drop_last=False,
                        num_workers=self.args.num_workers,
                        # persistent_workers=True,
                        worker_init_fn=worker_init_fn,
                        # pin_memory=True
                        )
                else:
                    sampler = dgl.dataloading.ShaDowKHopSampler(grad_block)
                    graph_loader = dgl.dataloading.DataLoader(
                        graph=batched_graph, graph_sampler=sampler,
                        batch_size=self.args.kernel_size,
                        shuffle=False,
                        drop_last=False,
                        num_workers=self.args.num_workers,
                        # persistent_workers=True,
                        worker_init_fn=worker_init_fn,
                        # pin_memory=True
                        )
                
                # 将每张图的 DataLoader 添加到列表
                self.graph_loaders.append(graph_loader)
                
                if self.args.frozen_padding >= 0:
                    self.graph_loaders_sec.append(graph_loader_sec)
            else:
                # TODO: 对整个图进行采样
                self.whole_graph = True
                # 你可以在这里为单个大图使用DataLoader
                pass
    
    @lru_cache(8)
    def id_in_parent(self, parent, sub):
        '''index transformation in subset'''
        if self.whole_graph:
            return sub
        sorted_parent, sorted_indices = torch.sort(parent)
        sorted_pos = torch.searchsorted(sorted_parent, sub)
        return sorted_indices[sorted_pos]
    
    def get_metrics(self, y_true, y_pred, y_proba, metrics=None):
        m = {}
        y_true = to_numpy(y_true)
        y_pred = to_numpy(y_pred)
        y_proba = to_numpy(y_proba)
        ytnan = np.isnan(y_true)
        ypnan = np.isnan(y_proba)
        nan_mask = ~(ytnan | ypnan)  # 找到所有非 NaN 的索引

        # 仅在发现 NaN 时才进行筛选，并记录 NaN 来源
        if not nan_mask.all():
            if ytnan.any():
                logger.info(f"Found {ytnan.sum()} NaN in y_true")
            if ypnan.any():
                logger.info(f"Found {ypnan.sum()} NaN in y_proba")

            y_true = y_true[nan_mask]
            y_proba = y_proba[nan_mask]
            
        for k, metric in metrics.items():
            m[k] = metric(y_true,y_pred,y_proba)
        return m
    
    # @torch.no_grad()
    # def evaluate(self, evaluator,epoch, e2e):
    #     # 设置模型为评估模式
    #     self.model_gnn.eval()
    #     if self.model_lm and e2e:
    #         self.model_lm.eval()
        
    #     # 初始化评估指标和损失
    #     preds = []
    #     # onehot_labels = self.cal_labels(self.n_node, self.labels)
    #     onehot_labels = self.cal_labels(self.n_node, [], None)
    #     if e2e:
    #         self.out_lm, self.feat_static = self.get_feat()
        
    #     # 遍历所有评估批次
    #     train_loss, val_loss, test_loss = 0.0, 0.0, 0.0
    #     n_train = len(self.data['train']['graphs'])
    #     n_val = len(self.data['valid']['graphs'])
    #     n_test = len(self.data['test']['graphs'])
    #     # with torch.no_grad():
        
    #     for batch in self.all_data_loader:  # 获取每个批次
    #         node_ids = batch['graph'].ndata['uid']  # 获取当前图的节点ID
    #         batched_graph = batch['graph'].to(device=self.device)  # 当前批次的DGL图
    #         node_type_feat = batch['graph'].ndata['type'].to(device=self.device)
    #         edge_type_feat = batch['graph'].edata['type'].to(device=self.device)
    #         # if e2e:
    #         #     feat_eval = self.get_feat(batch=batch, device=self.device)
    #         # else:
    #         feat_eval = self.feat_static[node_ids].to(self.device)  # 获取当前批次的特征
    #         out_lm = self.out_lm[node_ids].to(self.device)
    #         # loss = self.custom_loss(self.labels[node_ids], out_lm, feat = feat_eval, batch_size=256).item()
    #         # 如果使用标签，合并标签特征
    #         if self.args.use_labels:
    #             feat_eval = torch.cat([feat_eval, onehot_labels[node_ids]], dim=-1)

    #         # feat_eval = feat_eval.to(torch.float32)  # 确保数据类型正确
    #         feat_eval = torch.cat([node_type_feat, feat_eval], dim=-1)
    #         # 获取预测结果
    #         nnan = torch.isnan(feat_eval).sum()
    #         if nnan > 0:
    #             logger.warning(f"{torch.isnan(feat_eval).sum()} NaN in feat_eval")

    #         # 如果有标签迭代，进行n_label_iters次标签更新
    #         if self.args.n_label_iters > 0:
    #             gnn_pred = self.model_gnn(batched_graph, feat_eval, edge_type_feat)
    #             # unlabel_idx = torch.cat([torch.tensor(self.val_idx), torch.tensor(self.test_idx)])
    #             for _ in range(self.args.n_label_iters):
    #                 # gnn_pred = F.softmax(gnn_pred, dim=-1)
    #                 feat_eval[:,-self.args.num_labels:] = F.softmax(gnn_pred, dim=-1)
    #                 gnn_pred, gnn_feat = self.model_gnn(batched_graph, feat_eval, edge_type_feat, rt_feat = True)
    #         else:
    #             gnn_pred, gnn_feat = self.model_gnn(batched_graph, feat_eval, edge_type_feat, rt_feat = True)
                   
    #         pred = 0.5*(gnn_pred + out_lm)
    #         if(self.args.lm_x):
    #             pred = out_lm
    #         if(self.args.gnn_x):
    #             pred = gnn_pred
    #         # 存储预测结果和标签
    #         preds.append(pred)
    #     preds = torch.cat(preds, dim=0)
    #     # labels = torch.cat(labels, dim=0)
    #     y_probs = preds.softmax(dim=1)[:, 1]
    #     y_preds = (y_probs >= 0.4).long()
    #     if(self.args.lm_x):
    #             y_preds = (y_probs >= 0.6).long()
    #     # 计算损失
    #     train_metrics=[]
    #     val_metrics=[]
    #     test_metrics=[]
    #     if n_train:
    #         train_loss = self.criterion(preds[self.train_idx], self.labels[self.train_idx]).item()
    #         train_metrics = self.get_metrics(self.labels[self.train_idx], y_preds[self.train_idx], y_probs[self.train_idx], self.metrics)
        
    #     if n_val:
    #         val_loss = self.criterion(preds[self.val_idx], self.labels[self.val_idx]).item()
    #         val_metrics = self.get_metrics(self.labels[self.val_idx], y_preds[self.val_idx], y_probs[self.val_idx], self.metrics)
        
    #     if n_test:
    #         test_loss = self.criterion(preds[self.test_idx], self.labels[self.test_idx]).item()
    #         test_metrics = self.get_metrics(self.labels[self.test_idx], y_preds[self.test_idx], y_probs[self.test_idx],self.metrics_test)
    #     # 保存每个节点的预测结果
    #     if(self.args.test_only):
    #     # 构建 DataFrame
    #         node_results_df = pd.DataFrame({
    #             'node_id': list(range(len(y_preds))),  # 假设节点按顺序编号
    #             'y_true': self.labels.cpu().numpy(),
    #             'y_pred': y_preds.cpu().numpy(),
    #             'y_prob': y_probs.detach().cpu().numpy(),
    #         })

    #         # 创建输出目录（可选）
    #         output_dir = "./out/test_pred"
    #         os.makedirs(output_dir, exist_ok=True)

    #         # 构建文件路径，带上 epoch 信息便于追踪
    #         output_path = os.path.join(output_dir, f"cwe_node_preds.csv")

    #         # 保存为 CSV
    #         node_results_df.to_csv(output_path, index=False, encoding='utf-8-sig')
    #         logger.info(f"节点预测结果已保存到: {output_path}")


    #     torch.cuda.empty_cache()

    #     # 返回评估结果，包括准确率、损失和其他指标
    #     return  train_metrics, val_metrics, test_metrics, train_loss, val_loss, test_loss, preds
    @torch.no_grad()
    def evaluate(self, evaluator, epoch, e2e):

        # =====================================================
        # 0. 模式判断
        # =====================================================
        only_lm = self.args.lm_x
        only_gnn = self.args.gnn_x

        assert not (only_lm and only_gnn), \
            "错误：lm_x 和 gnn_x 不能同时为 True！"

        # =====================================================
        # 1. 设置 eval 模式
        # =====================================================
        self.model_gnn.eval()
        if self.model_lm:
            self.model_lm.eval()

        # =====================================================
        # 2. 如果 e2e，需要重新提取 LM 特征
        # =====================================================
        if e2e:
            self.out_lm, self.feat_static = self.get_feat()

        preds = []

        # =====================================================
        # 3. 遍历所有 batch
        # =====================================================
        for batch in self.all_data_loader:

            graph = batch["graph"].to(self.device)
            node_ids = graph.ndata["uid"]

            node_type_feat = graph.ndata["type"].to(self.device)
            edge_type_feat = graph.edata["type"].to(self.device)

            # =====================================================
            # ✅ Mode A: Only LM（只用语言模型）
            # =====================================================
            if only_lm:

                input_ids = batch["input_ids"].to(self.device)
                attention_mask = batch["attention_mask"].to(self.device)

                # eval 模式
                self.model_lm.eval()
                self.node_classifier.eval()

                # forward LM
                out_lm, lm_feat = self.model_lm(
                    input_ids,
                    attention_mask,
                    return_hidden=True
                )

                # CLS pooling
                if lm_feat.dim() == 3:
                    lm_feat = lm_feat[:, 0, :]

                # ✅node logits
                pred = self.node_classifier(lm_feat)

                preds.append(pred)
                continue


            # =====================================================
            # ✅ Mode B: Only GNN（只用图模型）
            # =====================================================
            if only_gnn:

                # frozen LM 特征作为输入
                feat_eval = self.feat_static[node_ids.cpu()].to(self.device)

                # 拼接节点类型
                feat_eval = torch.cat([node_type_feat, feat_eval], dim=-1)
                

                # forward GNN
                gnn_pred, gnn_feat = self.model_gnn(
                    graph, feat_eval, edge_type_feat, rt_feat=True
                )

                pred = gnn_pred
                preds.append(pred)
                continue   # ⚠️关键：跳过 LM fusion

            # =====================================================
            # ✅ Mode C: 融合模式（LM + GNN）
            # =====================================================

            # --- LM 输出 ---
            feat_eval = self.feat_static[node_ids.cpu()].to(self.device)
            out_lm = self.out_lm[node_ids.cpu()].to(self.device)

            feat_eval = torch.cat([node_type_feat, feat_eval], dim=-1)

            # --- forward GNN ---
            gnn_pred, gnn_feat = self.model_gnn(
                graph, feat_eval, edge_type_feat, rt_feat=True
            )

            # --- 融合预测 ---
            pred = 0.5 * (gnn_pred + out_lm)

            preds.append(pred)

        # =====================================================
        # 4. 拼接所有预测
        # =====================================================
        preds = torch.cat(preds, dim=0)

        # =====================================================
        # 5. 计算概率与分类结果
        # =====================================================
        y_probs = preds.softmax(dim=1)[:, 1]

        # 阈值可调
        threshold = 0.4

        y_preds = (y_probs >= threshold).long()

        # =====================================================
        # 6. 计算指标
        # =====================================================
        train_loss, val_loss, test_loss = 0.0, 0.0, 0.0
        train_metrics, val_metrics, test_metrics = [], [], []

        if len(self.train_idx) > 0:
            train_loss = self.criterion(
                preds[self.train_idx],
                self.labels[self.train_idx]
            ).item()

            train_metrics = self.get_metrics(
                self.labels[self.train_idx],
                y_preds[self.train_idx],
                y_probs[self.train_idx],
                self.metrics
            )

        if len(self.val_idx) > 0:
            val_loss = self.criterion(
                preds[self.val_idx],
                self.labels[self.val_idx]
            ).item()

            val_metrics = self.get_metrics(
                self.labels[self.val_idx],
                y_preds[self.val_idx],
                y_probs[self.val_idx],
                self.metrics
            )

        if len(self.test_idx) > 0:
            test_loss = self.criterion(
                preds[self.test_idx],
                self.labels[self.test_idx]
            ).item()

            test_metrics = self.get_metrics(
                self.labels[self.test_idx],
                y_preds[self.test_idx],
                y_probs[self.test_idx],
                self.metrics_test
            )

        # =====================================================
        # 7. test_only 保存结果
        # =====================================================
        if self.args.test_only:

            node_results_df = pd.DataFrame({
                "node_id": list(range(len(y_preds))),
                "y_true": self.labels.cpu().numpy(),
                "y_pred": y_preds.cpu().numpy(),
                "y_prob": y_probs.detach().cpu().numpy(),
            })

            output_dir = "./out/test_pred"
            os.makedirs(output_dir, exist_ok=True)

            output_path = os.path.join(output_dir, "cwe_node_preds.csv")
            node_results_df.to_csv(output_path, index=False, encoding="utf-8-sig")

            logger.info(f"节点预测结果已保存到: {output_path}")

        torch.cuda.empty_cache()

        return (
            train_metrics, val_metrics, test_metrics,
            train_loss, val_loss, test_loss,
            preds
        )

    

    
    def load_data(self):
        """加载多个图数据并根据训练集、验证集、测试集分割"""
        if self.args.test_only:
            #data_path = f"processed_data/{self.args.CWE_name}_all_dataset.pt"
            data_path = f"processed_data/{self.args.dataset}_all_dataset.pt"
        else:
            data_path = f"processed_data/{self.args.dataset}_all_dataset.pt"
        logger.info(f"Loading data from {data_path}...")
        
        try:
            # 从指定路径加载数据
            self.data = torch.load(data_path)
        except FileNotFoundError:
            logger.error(f"Data file not found at {data_path}")
            raise
        except Exception as e:
            logger.error(f"Error loading data from {data_path}: {e}")
            raise
        
        #self.evaluate = CodeEvaluator(self.args.dataset)

        # 提取训练集、验证集、测试集的数据
        self.graphs = self.data['train']['graphs'] + self.data['valid']['graphs'] + self.data['test']['graphs']
        self.text_tokens = self.data['train']['tokens'] + self.data['valid']['tokens'] + self.data['test']['tokens']

        self.labels = torch.cat([item['label'] for item in self.text_tokens], dim=0) 
        
        if not hasattr(self.args, "num_labels") or self.args.num_labels is None:
            self.args.num_labels = (self.labels.max() + 1).item()


        # print(self.labels.shape)
        self.uids = torch.cat([item['uid'] for item in self.text_tokens], dim=0)
       
        self.n_node=len(self.uids)
        logger.info(f"self.n_node: {self.n_node}")

        # 提取训练集、验证集、测试集的索引（基于 uid）
        self.train_idx = [uid for graph in self.data['train']['graphs'] for uid in graph.ndata['uid'].tolist()]
        self.val_idx = [uid for graph in self.data['valid']['graphs'] for uid in graph.ndata['uid'].tolist()]
        self.test_idx = [uid for graph in self.data['test']['graphs'] for uid in graph.ndata['uid'].tolist()]

        logger.info(f"Length of train_idx: {len(self.train_idx)}")
        logger.info(f"Length of val_idx: {len(self.val_idx)}")
        logger.info(f"Length of test_idx: {len(self.test_idx)}")
       
        # 检查数据是否为空
        if not self.graphs:
            logger.warning("Graph data is empty.")
        if not self.text_tokens:
            logger.warning("Token data is empty.")
        if self.labels is None:
            logger.warning("Label data is empty.")

        # 检查索引是否有效
        if self.train_idx is None or self.val_idx is None or self.test_idx is None:
            logger.warning("One or more index sets (train, val, test) are empty.")

        logger.info(f"Loaded {len(self.graphs)} graphs and {len(self.text_tokens)} token sets.")
        logger.info(f"Train indices: {len(self.train_idx)}, Val indices: {len(self.val_idx)}, Test indices: {len(self.test_idx)}")
        # print("Data loaded successfully.")
        # 如果需要，可以在这里进行进一步的数据处理或验证
        self._validate_data()
        
        self.args.n_node_feats = self.args.hidden_size
        #self.args.n_edge_types = 0 #从数据集读取
        self.args.n_edge_types = self.data['attr']['n_edge_types']

        if self.args.test_only:
            if not hasattr(self.args, "n_node_types") or self.args.n_node_types is None:
                # 这里 20 可以改为你训练时的维度，也可以通过参数传入
                self.args.n_node_types = 20
            logger.info(f"[Test-Only] Fix n_node_types = {self.args.n_node_types}")
        else:
            self.args.n_node_types = self.data['attr']['n_node_types']

        # ✅必须在 preprocess 前调用
        self.fix_test_node_type_dim()
        
        #self.text_token = TextDataset(input_ids, attention_mask, self.labels)
        self.evaluator = CodeEvaluator(name=self.args.dataset)
        self.split_idx = {}
        self.split_idx["train"] = self.train_idx
        self.split_idx["valid"] = self.val_idx
        self.split_idx["test"] = self.test_idx
        if self.args.use_labels:
            self.args.n_node_feats += self.args.num_labels
        self.args.n_node_feats += self.args.n_node_types
        self.feat = torch.empty((self.n_node, self.args.n_node_feats), dtype=torch.float32, device=self.device)
        # 确保将 train_idx 转换为张量
        self.train_idx = torch.tensor(self.train_idx, device=self.device)
        self.val_idx = torch.tensor(self.val_idx, device=self.device)
        self.test_idx = torch.tensor(self.test_idx, device=self.device)

        self.uid2text_index = {}
        for text_idx, token in enumerate(self.text_tokens):
            for uid in token['uid']:
                self.uid2text_index[uid.item()] = text_idx

        
        self.preprocess()
        all_dataset = GraphTextDataset(self.graphs, self.text_tokens, tokenizer=self.tokenizer)
        train_dataset = GraphTextDataset(self.graphs[:len(self.data['train']['graphs'])], self.data['train']['tokens'], tokenizer=self.tokenizer)
        self.train_data_loader = DataLoader(train_dataset, batch_size=self.args.batch_size_train, shuffle=False, collate_fn=collate_fn,
                        num_workers=self.args.num_workers,
                        # persistent_workers=True,
                        worker_init_fn=worker_init_fn,
                        # pin_memory=True
                        )#True
        self.all_data_loader = DataLoader(all_dataset, batch_size=self.args.batch_size_train, shuffle=False, collate_fn=collate_fn,
                        num_workers=self.args.num_workers,
                        # persistent_workers=True,
                        worker_init_fn=worker_init_fn,
                        # pin_memory=True
                        )#True
       
        input_ids_list = []
        attention_mask_list = []

        for batch in self.all_data_loader:
            input_ids_list.append(batch['input_ids'])
            attention_mask_list.append(batch['attention_mask'])

        input_ids = torch.cat(input_ids_list, dim=0)
        attention_mask = torch.cat(attention_mask_list, dim=0)

        self.text_token = TextDataset(input_ids, attention_mask, self.labels)

        del input_ids, attention_mask

    def fix_test_node_type_dim(self):
        """
        ✅只在 test_only 时运行：
        自动 pad/crop node_type_feat，使其维度与训练模型一致
        """

        if not self.args.test_only:
            return

        target_dim = self.args.n_node_types

        test_graphs = self.data["test"]["graphs"]

        for i, g in enumerate(test_graphs):

            if "type" not in g.ndata:
                continue

            feat = g.ndata["type"]
            cur_dim = feat.shape[1]

            # ---------- Pad ----------
            if cur_dim < target_dim:
                pad = torch.zeros(
                    feat.shape[0],
                    target_dim - cur_dim,
                    device=feat.device
                )
                feat = torch.cat([feat, pad], dim=1)


            # ---------- Crop ----------
            elif cur_dim > target_dim:
                feat = feat[:, :target_dim]

                logger.info(
                    f"[Fix-Type] Graph {i}: crop {cur_dim} → {target_dim}"
                )

            g.ndata["type"] = feat


    def _validate_data(self):
        """验证加载的数据是否有效"""
        if len(self.graphs) != len(self.text_tokens):
            logger.error("Mismatch between number of graphs and tokens.")
            raise ValueError("Mismatch between number of graphs and tokens.")
        if len(self.labels) != len(self.uids):
            logger.error("Mismatch between number of labels and uids.")
            raise ValueError("Mismatch between number of labels and uids.")
        if max(self.train_idx + self.val_idx + self.test_idx) > max(self.uids):
            logger.error("Index out of range: some indices exceed the maximum uid.")
            raise ValueError("Index out of range: some indices exceed the maximum uid.")
        max_uid = max(self.uids)
        max_index = max(self.train_idx + self.val_idx + self.test_idx)

        if max_index > max_uid:
            # 打印超出范围的索引
            invalid_train_indices = [idx for idx in self.train_idx if idx > max_uid]
            invalid_val_indices = [idx for idx in self.val_idx if idx > max_uid]
            invalid_test_indices = [idx for idx in self.test_idx if idx > max_uid]

            if invalid_train_indices:
                print("Invalid train indices (exceed max UID):", invalid_train_indices)
            if invalid_val_indices:
                print("Invalid validation indices (exceed max UID):", invalid_val_indices)
            if invalid_test_indices:
                print("Invalid test indices (exceed max UID):", invalid_test_indices)

            logger.error(f"Index out of range: max index ({max_index}) exceeds max UID ({max_uid}).")
            raise ValueError(f"Index out of range: max index ({max_index}) exceeds max UID ({max_uid}).")
        else:
            logger.info(f"Index validation passed: max index ({max_index}) <= max UID ({max_uid}).")



    def get_graph_loader(self, batch_size):
        """ 生成图数据加载器 """
        pass
    
    def get_params(self, init=False, custom_lr = False, need_name = False, gm_lora = False, grad_only = True):
        params = []
        if init:
            self.is_lm = []
            if self.model_lm:
                lmp = [{'params': p, 'lr': self.args.lm_lr} for p in self.model_lm.parameters()] if custom_lr else\
                    list(self.model_lm.parameters())
                params += lmp
                self.is_lm += [1 for _ in range(len(lmp))]
            if self.model_gnn:
                gmp = [{'params': p, 'lr': self.args.gm_lr} for p in self.model_gnn.parameters()]if custom_lr else\
                    list(self.model_gnn.parameters())
                params += gmp
                self.is_lm += [0 for _ in range(len(gmp))]
            if gm_lora: 
                params = []
                for n, p in self.model_gnn.named_parameters():
                    if 'lora_' in n:
                        if custom_lr:
                            params.append({'params': p, 'lr': self.args.gm_lr})
                        else:
                            params.append(p)
                return params
            if grad_only:
                if custom_lr:
                    return [p for p in params if p['params'].requires_grad]
                else:
                    return [p for p in params if p.requires_grad]
        elif need_name:
            if self.model_lm:
                params += list(self.model_lm.named_parameters())
            if self.model_gnn:
                params += list(self.model_gnn.named_parameters())
            if grad_only:
                return [(n,p) for (n,p) in params if p.requires_grad]
        else:
            if self.model_lm:
                params += list(self.model_lm.parameters())
            if self.model_gnn:
                params += list(self.model_gnn.parameters())
            if grad_only:
                return [p for p in params if p.requires_grad]
        return params

    def init_gnn(self):
        if self.args.gnn_type == "RevGAT":
            self.model_gnn = RevGAT(
                self.args,
                activation=F.leaky_relu,
                dropout=self.args.dropout,
                input_drop=self.args.input_drop,
                attn_drop=self.args.attn_drop,
                edge_drop=self.args.edge_drop,
                use_attn_dst=not self.args.no_attn_dst,
                use_symmetric_norm=self.args.use_norm,
                use_gpt_preds=self.args.use_gpt_preds,
                se_reduction = self.args.se_reduction
            )

            if self.args.ckpt_dir != '' and os.path.exists(self.args.ckpt_dir):
                self.model_gnn.load_state_dict(torch.load(self.args.ckpt_dir),strict=False)
                logger.info(f"Loaded PGM from {self.args.ckpt_dir}")
                self.model_gnn.convs[-1].reset_parameters()
                
        elif self.args.gnn_type == "GraphSAGE":
            print(self.args.n_node_feats)
            self.model_gnn = HSAGE(
                in_channels=self.args.n_node_feats,
                hidden_channels=self.args.n_hidden,
                out_channels=self.args.num_labels,
                num_layers=self.args.n_layers,
                dropout=self.args.dropout,
                # use_gpt_preds=self.args.use_gpt_preds
            )
        else:
            raise Exception(f"Unknown gnn {self.args.gnn_type}")
        
        self.optimizer = optim.RMSprop(self.get_params(init=True, custom_lr=True, grad_only=False), lr=self.args.gm_lr, weight_decay=self.args.wd)
        self.require_grad = [0 for _ in self.model_lm.parameters()]
        self.get_lr()
        # self.scheduler = lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=int((self.args.n_epochs-self.args.warmup)), eta_min=1e-7)
        return 1
    
    def init_lm(self):
        if not self.args.use_external_feat:
            self.model_lm = get_model_class(self.args.lm_type,self.args.task_type)(self.args)
            if hasattr(self.model_lm, 'tokenize'):
                self.tokenizer = self.model_lm.tokenize
            else:
                self.tokenizer = self.args.pretrained_repo
                os.environ["TOKENIZERS_PARALLELISM"] = "false"

        
        return 
    
    def init_optimizer_lm(self):
        """Only LM 模式使用"""

        if self.model_lm is None:
            return

        self.optimizer_lm = torch.optim.AdamW(
            self.model_lm.parameters(),
            lr=self.args.lm_lr
        )

        logger.info("LM optimizer 初始化完成")

    def train_lm(self):
        # trainer
        # Trainer = get_trainer_class(self.args.lm_type,self.args.dataset)
        trainer = Code_Trainer(self.args, self.text_token, self.split_idx, self.evaluator, self.model_lm)
        results, x_embs, self.model_lm = trainer.train()
        del trainer, results, x_embs
        torch.cuda.empty_cache()
    
    def load_stat(self):
        out_dir = f"{self.args.save}/ckpt"
        fname = os.path.join(out_dir, f"last_stat.pt")
        checkpoint = torch.load(fname, map_location=self.device)
        last_epoch = checkpoint['epoch']  # 从上次结束的epoch开始
        e2e = checkpoint['full_ft'] 
        best = 0.77
        logger.info(f"Loading last ckpt from {fname}, continue after ep{last_epoch}")
        if 0 < self.args.peft_start <= last_epoch:
            self.switch_to('gnn_lora')     
        self.to_device(self.model_gnn)
        self.to_device(self.model_lm)
        self.model_gnn.load_state_dict(checkpoint['gnn_dict'],strict=False)
        if self.model_lm: self.model_lm.load_state_dict(checkpoint['lm_dict'])
        # if self.args.peft_start != last_epoch: 
        self.optimizer.load_state_dict(checkpoint['optm_dict'])
        self.feat_static = checkpoint['feat_static']
        # self.args = checkpoint['args']
        # best = checkpoint['best']
        del checkpoint
        # logger.info(f"Loaded args: {self.args}")
        return last_epoch, e2e, best
    
    def switch_to(self, mode: Literal['gnn_lora', 'gnn_backbone', 'gnn_only', 'e2e','lm_only']):
        if mode=='gnn_lora':
            if not self.lora_added:
                src_model = self.model_gnn.state_dict()
                # src_model.to('cpu')
                self.model_gnn = RevGAT(
                    self.args,
                    activation=F.relu,
                    dropout=self.args.dropout,
                    input_drop=self.args.input_drop,
                    attn_drop=self.args.attn_drop,
                    edge_drop=self.args.edge_drop,
                    use_attn_dst=not self.args.no_attn_dst,
                    use_symmetric_norm=self.args.use_norm,
                    use_gpt_preds=self.args.use_gpt_preds,
                    lora_params={
                        'use_lora': self.args.use_peft,
                        'r': self.args.peft_r_gm,
                        'lora_alpha': self.args.peft_lora_alpha,
                        'lora_dropout': self.args.peft_lora_dropout
                        }
                    )
                self.model_gnn.load_state_dict(src_model, strict=False)
                self.to_device(self.model_gnn)
                self.lora_added = True
                # self.args.gm_lr = self.optimizer.param_groups[-1]["lr"]
                for item in self.get_params(init=True, custom_lr=True, gm_lora=True):
                    self.optimizer.add_param_group(item)

            lora.mark_only_lora_as_trainable(self.model_gnn)
            logger.info("GM switched to LoRA")
        elif mode == 'gnn_backbone':
            lora.mark_not_lora_as_trainable(self.model_gnn)
            logger.info("GM switched to backbone")
        elif mode == 'gnn_only':
            for i, p in enumerate(self.model_lm.parameters()):
                if p.requires_grad == True:
                    p.requires_grad = False
                    self.require_grad[i] = 1
            logger.info("Switched to gnn only")
        elif mode == 'e2e':
            for i, p in enumerate(self.model_lm.parameters()):
                if self.require_grad[i]:
                    p.requires_grad = True
            logger.info("Switched to lm+gnn")

        else:
            logger.info(f"Invalid switch mode: {mode}")
        logger.info(f"Number of trainable params: {self.count_params()}")
        
    def get_lr(self, eta_min=1e-3):
        self.lm_lr = [0]
        self.gm_lr = [0]
        total_epochs = len(self.args.ftmask)
        decay_epochs = total_epochs - self.args.warmup
        if decay_epochs > 0:
            k_lm = np.log(eta_min*self.args.gm_lr*0.3 / self.args.lm_lr) / decay_epochs
            k_gm = np.log(eta_min) / decay_epochs
        else:
            k_lm = k_gm = 0 
        
        for ep, e2e in enumerate(self.args.ftmask):
            if ep == 0:
                continue
            if ep <= self.args.warmup:
                lm_lr = self.args.lm_lr * ep / self.args.warmup
                gm_lr = self.args.gm_lr * ep / self.args.warmup
            else:
                decay_step = ep - self.args.warmup
                lm_lr = self.args.lm_lr * np.exp(k_lm * decay_step)
                gm_lr = self.args.gm_lr * np.exp(k_gm * decay_step)

            if e2e > 0:
                gm_lr /= 800
                lm_lr /= 800
            
            self.lm_lr.append(lm_lr)
            self.gm_lr.append(gm_lr)

            
    def adjust_learning_rate(self, epoch):
        '''lr schedule'''
            
        for i, param_group in enumerate(self.optimizer.param_groups):
            if self.is_lm[i]:
                param_group["lr"] = self.lm_lr[epoch]
            else:
                param_group["lr"] = self.gm_lr[epoch]

        logger.info(f"lm_lr: {self.lm_lr[epoch]}, gm_lr: {self.gm_lr[epoch]}")

    @torch.no_grad()
    def compute_tam(self, output, edge_index, label, train_mask, class_num_list=None, temp_phi=None, temp_gamma=None):
        eps = 1e-8
        n_cls = 2  # 假设是二分类问题
        
        # 1. 验证输入数据
        assert torch.all(label[train_mask] >= 0) and torch.all(label[train_mask] < n_cls), "无效的标签值"
        assert train_mask.sum() > 0, "train_mask不能为空"
        
        # 2. 温度缩放和邻居聚合
        #cls_num_list = torch.FloatTensor(class_num_list).to(output.device)
        cls_num_list = class_num_list.float()
        cls_num_ratio = cls_num_list / cls_num_list.sum()
        cls_num_ratio = cls_num_ratio * temp_gamma + (1 - temp_gamma)
        max_beta = torch.max(cls_num_ratio)
        cls_temperature = (temp_phi * (cls_num_ratio + 1 - max_beta)).unsqueeze(0)
        temp = 1 / cls_temperature.clamp(min=eps)

        # 3. 计算邻居分布
        agg_out = F.softmax(output.clone().detach()/temp, dim=1)
        agg_out[train_mask] = F.one_hot(label[train_mask].clone(), num_classes=n_cls).float()
        
        aggregator = MeanAggregation()
        neighbor_dist = aggregator(agg_out, edge_index)[train_mask]
        neighbor_dist = neighbor_dist.clamp(min=eps, max=1-eps)  # 严格限制在(0,1)范围内

        # 4. 计算connectivity_matrix（增强鲁棒性）
        connectivity_matrix = []
        for c in range(n_cls):
            c_mask = (label[train_mask] == c)
            if c_mask.sum() == 0:
                # 如果该类别没有样本，使用均匀分布
                connectivity_matrix.append(torch.ones_like(neighbor_dist[0]) / n_cls)
            else:
                # 添加平滑处理
                class_mean = neighbor_dist[c_mask].mean(dim=0)
                class_mean = class_mean.clamp(min=eps, max=1-eps)
                class_mean = class_mean / class_mean.sum()  # 重新归一化
                connectivity_matrix.append(class_mean)
        
        connectivity_matrix = torch.stack(connectivity_matrix, dim=0)
        connectivity_matrix = connectivity_matrix.clamp(min=eps, max=1-eps)

        # 5. 计算ACM（保持不变）
        center_mask = F.one_hot(label[train_mask].clone(), num_classes=n_cls).bool()
        acm = (neighbor_dist[center_mask].unsqueeze(dim=1) / torch.diagonal(connectivity_matrix).unsqueeze(dim=1)[label[train_mask]]) \
                    * (connectivity_matrix[label[train_mask]] / neighbor_dist)
        acm = acm.clamp(min=0, max=1)
        acm[center_mask] = 1

        # 6. 计算ADM（全面稳定性增强）
        # 6.1 计算JSD距离
        def safe_jsd(p, q):
            p = p.clamp(min=eps, max=1-eps)
            q = q.clamp(min=eps, max=1-eps)
            m = 0.5 * (p + q)
            return 0.5 * F.kl_div(torch.log(p), m, reduction='none') + 0.5 * F.kl_div(torch.log(q), m, reduction='none')

        # 6.2 计算各项分量
        cls_pair_jsd = safe_jsd(connectivity_matrix.unsqueeze(0), 
                            connectivity_matrix.unsqueeze(1)).sum(dim=-1)
        cls_pair_jsd = cls_pair_jsd.clamp(min=eps)
        
        self_kl = safe_jsd(neighbor_dist, 
                        connectivity_matrix[label[train_mask]]).sum(dim=-1, keepdim=True)
        neighbor_kl = safe_jsd(neighbor_dist.unsqueeze(1),
                            connectivity_matrix.unsqueeze(0)).sum(dim=-1)

        # 6.3 稳定计算ADM
        numerator = (self_kl**2 + (cls_pair_jsd**2)[label[train_mask]] - neighbor_kl**2)
        numerator = numerator.clamp(min=-1e6, max=1e6)  # 防止极端值
        
        denominator = (2 * (cls_pair_jsd**2)[label[train_mask]]).clamp(min=eps, max=1e6)
        adm = numerator / denominator
        
        # 6.4 处理异常情况
        adm = torch.nan_to_num(adm, nan=0.0, posinf=0.0, neginf=0.0)
        adm[center_mask] = 0
        
        # 最终验证
        assert not torch.isnan(acm).any(), "ACM中有NaN"
        assert not torch.isnan(adm).any(), "ADM中有NaN"
        
        return acm, adm
    
    def adjust_output(self, output, edge_index, label, train_mask, class_num_list, epoch):
        if self.args.use_tam and epoch > self.args.warmup_tam:
            # 检查输入
            assert not torch.isnan(output).any(), "NaN in output before TAM"
            assert train_mask.any(), "train_mask is empty in adjust_output!"

            # 计算 ACM 和 ADM
            acm, adm = self.compute_tam(
                output=output,
                edge_index=edge_index,
                label=label,
                train_mask=train_mask,
                class_num_list=class_num_list,
                temp_phi=self.args.temp_phi,
                temp_gamma=0.4  # 或其他默认值
            )
            
            # 检查 acm/adm 是否有效
            assert not torch.isnan(acm).any(), "NaN in ACM"
            assert not torch.isnan(adm).any(), "NaN in ADM"

            # 调整输出（限制幅度避免爆炸）
            acm = acm.clamp(min=-10, max=10).log()
            adm = adm.clamp(min=-10, max=10)
            output[train_mask] = output[train_mask] + self.args.tam_alpha * acm + self.args.tam_beta * adm

            # 检查输出是否有效
            assert not torch.isnan(output).any(), "NaN in output after TAM"
        return output

    def get_feat(self, batch=None, device='cpu', return_cls=True):
        # 如果传入了batch，直接使用传入的batch
        if batch is not None:
            if self.args.batch_size_infer < len(batch['input_ids']):
                text_data = TensorDataset(batch['input_ids'],batch['attention_mask'])
            else:
                input_ids = batch['input_ids'].to(self.device)
                attention_mask = batch['attention_mask'].to(self.device)
                # if self.args.fp16:
                #     with autocast(): 
                        # out, feat = self.model_lm(input_ids, attention_mask, return_hidden=True)
                # else:
                out, feat = self.model_lm(input_ids, attention_mask, return_hidden=True)
                if return_cls:
                    return out, feat.to(device)
                else:
                    return feat.to(device)
        else:
            #NOTE：这里写直接用text_token 获取全部feat
            text_data = TensorDataset(self.text_token.get_value('input_ids'), self.text_token.get_value('att_mask'))

        n_size = len(text_data) 
        # 创建feat和out矩阵，矩阵的大小根据当前批次大小动态调整
        batch_size = self.args.batch_size_infer
        text_loader = DataLoader(text_data, batch_size=batch_size, shuffle=False,
                        num_workers=self.args.num_workers,
                        # persistent_workers=True,
                        worker_init_fn=worker_init_fn,
                        # pin_memory=True
                        ) 
        num_batches = len(text_loader)
        interval = num_batches // 10
        
        feat = torch.empty((n_size, self.args.hidden_size), dtype=self.dtype, device=self.device)
        out = torch.empty((n_size, self.args.num_labels), dtype=self.dtype, device=self.device)

        with tqdm(total=num_batches, desc='LM inference', unit='batch', file=open(os.devnull, 'w')) as pbar:
            for i, (input_ids, attention_mask) in enumerate(text_loader):

                input_ids = input_ids.to(self.device)  # 将数据移到设备上
                attention_mask = attention_mask.to(self.device)

                # 使用FP16推理（如果指定）
                # if self.args.fp16:
                #     with autocast():    
                #         out[i*batch_size:i*batch_size+input_ids.shape[0]], \
                #         feat[i*batch_size:i*batch_size+input_ids.shape[0],:self.args.hidden_size] = self.model_lm(input_ids, attention_mask, return_hidden=True)
                # else:
                out[i*batch_size:i*batch_size+input_ids.shape[0]], \
                feat[i*batch_size:i*batch_size+input_ids.shape[0],:self.args.hidden_size] = self.model_lm(input_ids, attention_mask, return_hidden=True)

                # 更新进度条
                if interval == 0: 
                    logger.info(str(pbar))
                elif (i-1) % interval == 0: 
                    logger.info(str(pbar))

                pbar.update(1)
                
        # if out_batch.shape[0] == 0 or feat_batch.shape[0] == 0:
        #     raise ValueError("Empty output from self.model_lm detected!")
        if return_cls:
            return out, feat.to(device)
        else:
            return feat.to(device)

        
    # def train(self, epoch, evaluator, e2e, mode="teacher", teacher_output=None):
    #     self.model_gnn.train()
    #     if self.model_lm and e2e:
    #         self.model_lm.train()

    #     # ---------- 预计算 LM 特征 ----------
    #     if self.feat_static is None:
    #         with torch.no_grad():
    #             self.out_lm, self.feat_static = self.get_feat()

    #     num_batches = len(self.train_data_loader)
    #     interval = max(1, num_batches // 10)

    #     with tqdm(total=num_batches,
    #             desc=f"Train {epoch}/{self.args.n_epochs}",
    #             unit="batch",
    #             file=open(os.devnull, 'w')) as pbar:

    #         for step, batch in enumerate(self.train_data_loader):

    #             # ===== 1. 取 batch 数据 =====
    #             graph = batch["graph"].to(self.device)
    #             node_ids = graph.ndata["uid"]
    #             labels = batch["labels"].to(self.device)

    #             input_ids = batch["input_ids"].to(self.device)
    #             attention_mask = batch["attention_mask"].to(self.device)

    #             node_type_feat = graph.ndata["type"]
    #             edge_type_feat = graph.edata["type"]

    #             num_nodes = graph.num_nodes()

    #             # ===== 2. LM 特征 =====
    #             if e2e and (num_nodes <= self.args.grad_size or self.whole_graph):
    #                 out_lm, lm_feat = self.model_lm(
    #                     input_ids, attention_mask, return_hidden=True
    #                 )
    #             else:
    #                 lm_feat = self.feat_static[node_ids.cpu()].to(self.device)
    #                 out_lm = self.out_lm[node_ids.cpu()].to(self.device)

    #             feat = lm_feat

    #             # ===== 3. 拼接节点类型 =====
    #             feat = torch.cat([node_type_feat, feat], dim=-1)

    #             # ===== 4. Label masking =====
    #             if self.args.use_labels:
    #                 onehot = self.cal_labels(num_nodes, labels)
    #                 feat = torch.cat(
    #                     [feat, torch.zeros(num_nodes, self.args.num_labels,
    #                                     device=self.device)],
    #                     dim=-1
    #                 )

    #                 perm = torch.randperm(num_nodes, device=self.device)
    #                 n_pred = max(1, int(num_nodes * self.args.mask_rate))
    #                 train_pred_idx = node_ids[perm[:n_pred]]
    #                 train_label_idx = node_ids[perm[n_pred:]]

    #                 if len(train_label_idx) > 0:
    #                     feat[self.id_in_parent(node_ids, train_label_idx),
    #                         -self.args.num_labels:] = \
    #                         onehot[self.id_in_parent(node_ids, train_label_idx)]
    #             else:
    #                 mask = torch.rand(num_nodes, device=self.device) < self.args.mask_rate
    #                 train_pred_idx = node_ids[mask]

    #             assert train_pred_idx.numel() > 0

    #             # ===== 5. GNN Forward =====
    #             self.optimizer.zero_grad()

    #             pred, gnn_feat = self.model_gnn(
    #                 graph, feat, edge_type_feat, rt_feat=True
    #             )

    #             # ===== 6. LM-GNN 融合（LM 主导）=====
    #             alpha = getattr(self.args, "residual_alpha", 0.3)
    #             fused_feat = gnn_feat + alpha * lm_feat
    #             if(self.args.lm_x):
    #                 fused_feat = lm_feat

    #             if(self.args.gnn_x):
    #                 fused_feat = gnn_feat
    #             # ===== 7. 只在 mask 节点上算 loss =====
    #             local_idx = self.id_in_parent(node_ids, train_pred_idx)

    #             pred_m = pred[local_idx]
    #             labels_m = labels[local_idx]

    #             loss = 0.0

    #             if e2e:
    #                 loss += self.args.lm_loss_weight * self.custom_loss(
    #                     labels_m,
    #                     out_lm[local_idx],
    #                     lm_feat[local_idx]
    #                 )

    #                 loss += (1 - self.args.lm_loss_weight) * self.custom_loss(
    #                     labels_m,
    #                     pred_m,
    #                     fused_feat[local_idx]
    #                 )
    #             else:
    #                 loss = self.custom_loss(
    #                     labels_m,
    #                     pred_m,
    #                     fused_feat[local_idx]
    #                 )

    #             # ===== 8. 反向传播 =====
    #             if self.args.fp16:
    #                 self.scaler.scale(loss).backward()
    #                 self.scaler.step(self.optimizer)
    #                 self.scaler.update()
    #             else:
    #                 loss.backward()
    #                 self.optimizer.step()

    #             # ===== 9. 日志 =====
    #             if step % interval == 0:
    #                 logger.info(f"[Epoch {epoch}] step {step}/{num_batches}, loss={loss.item():.4f}")

    #             pbar.update(1)

    #         torch.cuda.empty_cache()

    def train(self, epoch, evaluator, e2e, mode="teacher", teacher_output=None):

        # =====================================================
        # 0. 模式判断
        # =====================================================
        only_lm = self.args.lm_x
        only_gnn = self.args.gnn_x

        assert not (only_lm and only_gnn), \
            "错误：lm_x 和 gnn_x 不能同时为 True！"

        # =====================================================
        # 1. 设置模型状态（train / eval）
        # =====================================================
        if only_lm:
            # Only LM：只训练语言模型
            self.model_lm.train()
            self.model_gnn.eval()

        elif only_gnn:
            # Only GNN：只训练图模型
            self.model_gnn.train()
            self.model_lm.eval()

        else:
            # 联合训练
            self.model_gnn.train()
            if self.model_lm and e2e:
                self.model_lm.train()
            else:
                self.model_lm.eval()

        # =====================================================
        # 2. 预计算 LM 静态特征（Only GNN / 非E2E需要）
        # =====================================================
        if self.feat_static is None:
            with torch.no_grad():
                self.out_lm, self.feat_static = self.get_feat()

        num_batches = len(self.train_data_loader)
        interval = max(1, num_batches // 10)

        # =====================================================
        # 3. Train Loop
        # =====================================================
        for step, batch in enumerate(self.train_data_loader):

            # -------------------------------------------------
            # (1) 取 batch 数据
            # -------------------------------------------------
            graph = batch["graph"].to(self.device)
            node_ids = graph.ndata["uid"]
            labels = batch["labels"].to(self.device)

            input_ids = batch["input_ids"].to(self.device)
            attention_mask = batch["attention_mask"].to(self.device)

            node_type_feat = graph.ndata["type"]
            edge_type_feat = graph.edata["type"]

            num_nodes = graph.num_nodes()

            # -------------------------------------------------
            # (2) Mask 节点（只在部分节点算 loss）
            # -------------------------------------------------
            perm = torch.randperm(num_nodes, device=self.device)

            n_pred = max(1, int(num_nodes * self.args.mask_rate))

            train_pred_idx = node_ids[perm[:n_pred]]


            local_idx = self.id_in_parent(node_ids, train_pred_idx)
            labels_m = labels[local_idx]

            # =====================================================
            # ✅ Mode A: Only LM（只训练语言模型）
            # =====================================================
            if only_lm:

                # 1）确保 LM 参数可训练
                for name, p in self.model_lm.named_parameters():
                    if "lora" not in name:
                        p.requires_grad = False

                out_lm, lm_feat = self.model_lm(
                    input_ids,
                    attention_mask,
                    return_hidden=True
                )
                

                if lm_feat.dim() == 3:
                    lm_feat = lm_feat[:, 0, :]   # CLS pooling

                node_logits = self.node_classifier(lm_feat)

                pred_m = node_logits[local_idx]

                loss = self.criterion(pred_m, labels_m)

                self.optimizer_lm.zero_grad()
                loss.backward()
                self.optimizer_lm.step()



            # =====================================================
            # ✅ Mode B: Only GNN（只训练图模型）
            # =====================================================
            elif only_gnn:

                # --- LM 特征 frozen ---
                with torch.no_grad():
                    lm_feat = self.feat_static[node_ids.cpu()].to(self.device)

                # --- 输入特征 = 节点类型 + LM embedding ---
                feat = torch.cat([node_type_feat, lm_feat], dim=-1)

                if feat.shape[1] != self.args.n_node_feats:
                    pad_dim = self.args.n_node_feats - feat.shape[1]
                    feat = torch.cat(
                        [feat, torch.zeros(num_nodes, pad_dim, device=self.device)],
                        dim=-1
                    )

                # --- forward GNN ---
                pred, gnn_feat = self.model_gnn(
                    graph,
                    feat,
                    edge_type_feat,
                    rt_feat=True
                )

                pred_m = pred[local_idx]

                # --- loss ---
                loss = self.custom_loss(
                    labels_m,
                    pred_m,
                    gnn_feat[local_idx]
                )

                # --- backward：只更新 GNN ---
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

            # =====================================================
            # ✅ Mode C: LM + GNN 联合训练（融合）
            # =====================================================
            else:

                # --- LM 特征 ---
                if e2e:
                    out_lm, lm_feat = self.model_lm(
                        input_ids,
                        attention_mask,
                        return_hidden=True
                    )
                else:
                    lm_feat = self.feat_static[node_ids.cpu()].to(self.device)
                    out_lm = self.out_lm[node_ids.cpu()].to(self.device)

                # --- 拼接节点类型 ---
                feat = torch.cat([node_type_feat, lm_feat], dim=-1)

                # --- forward GNN ---
                pred, gnn_feat = self.model_gnn(
                    graph,
                    feat,
                    edge_type_feat,
                    rt_feat=True
                )

                # --- 特征融合 ---
                alpha = getattr(self.args, "residual_alpha", 0.3)
                fused_feat = gnn_feat + alpha * lm_feat

                pred_m = pred[local_idx]

                # --- loss（融合输出训练） ---
                loss = self.custom_loss(
                    labels_m,
                    pred_m,
                    fused_feat[local_idx]
                )

                # --- backward ---
                self.optimizer.zero_grad()
                loss.backward()
                self.optimizer.step()

            # =====================================================
            # 4. Logging
            # =====================================================
            if step % interval == 0:
                logger.info(
                    f"[Epoch {epoch}] Step {step}/{num_batches}, "
                    f"Loss={loss.item():.4f}"
                )

        torch.cuda.empty_cache()

    


    def load_model(self, model_path):
        if self.args.test_only:
            logger.info(f"从 {model_path} 加载模型...")
            checkpoint = torch.load(model_path, map_location=self.device)

            # 这段逻辑需要和你的 load_data 函数中计算 n_node_feats 的逻辑保持完全一致
            self.args.n_node_feats = self.args.hidden_size 
            if self.args.use_labels:
                self.args.n_node_feats += self.args.num_labels
            self.args.n_node_feats += self.args.n_node_types # 这里会使用上面我们手动设置的正确值
            # --- 关键修改结束 ---


            # 现在，用被修正过的正确维度参数来创建模型
            self.init_lm()
            self.init_gnn() 

            self.model_gnn.load_state_dict(checkpoint['gnn_dict'], strict=False)
            # 使用 .get() 方法安全访问，避免 KeyError
            if self.model_lm and checkpoint.get('lm_dict'):
                self.model_lm.load_state_dict(checkpoint['lm_dict'])
            
            if 'feat_static' in checkpoint:
                self.feat_static = checkpoint['feat_static']
            
            self.to_device(self.model_gnn)
            self.to_device(self.model_lm)
            logger.info("模型加载成功，准备评估。")
            
        else:
            """加载已保存的模型状态字典。"""
            logger.info(f"从 {model_path} 加载模型...")
            checkpoint = torch.load(model_path, map_location=self.device)

            # 确保模型实例在加载状态字典前被创建
            self.init_lm()  # 确保 LM 模型被初始化，因为 tokenizer 在这里设置
            self.init_gnn() # 确保 GNN 模型被初始化

            self.model_gnn.load_state_dict(checkpoint['gnn_dict'], strict=False)
            # 使用 .get() 方法安全访问，避免 KeyError
            if self.model_lm and checkpoint.get('lm_dict'):
                self.model_lm.load_state_dict(checkpoint['lm_dict'])
            
            if 'feat_static' in checkpoint:
                self.feat_static = checkpoint['feat_static']
            
            self.to_device(self.model_gnn)
            self.to_device(self.model_lm)
            logger.info("模型加载成功，准备评估。")
            

    def evaluate_standalone(self, model_path):
        """
        加载预训练模型，在新数据集上进行评估，并记录详细的性能指标。
        """
        logger.info(f"启动独立评估模式。")

        # --- 步骤 1: 加载数据和模型 ---
        self.load_data()
        self.prepare() 
        self.load_model(model_path)
        
        # --- 新增：在模型加载后立即计算参数数量 ---
        gnn_params = sum(p.numel() for p in self.model_gnn.parameters() if p.requires_grad)
        total_params = gnn_params
        is_e_2_e_eval = (self.model_lm is not None) and not self.args.use_external_feat
        if is_e_2_e_eval:
            lm_params = sum(p.numel() for p in self.model_lm.parameters() if p.requires_grad)
            total_params += lm_params
            with torch.no_grad():
                self.out_lm, self.feat_static = self.get_feat()
        else:
            if self.feat_static is None:
                logger.warning("非端到端模式下 feat_static 为 None，评估可能无法进行。")

        # --- 步骤 2: 准备评估并启动性能计时 ---
        evaluator_wrapper = lambda pred, labels: self.evaluator.eval(
            {"y_pred": pred.argmax(dim=-1, keepdim=True), "y_true": labels}
        )["acc"]

        logger.info("开始在新数据集上进行推理评估...")
        # --- 新增：性能计量准备 ---
        torch.cuda.reset_peak_memory_stats(self.device) # 重置显存峰值统计
        torch.cuda.synchronize() # 等待CUDA操作完成，确保计时准确
        start_time = time.time()

        # --- 步骤 3: 执行核心评估 ---
        # 调用 evaluate 函数，这一部分保持不变
        train_metrics, val_metrics, test_metrics, train_loss, val_loss, test_loss, pred = \
                self.evaluate(evaluator_wrapper, 0, True)

        # --- 新增：结束性能计时并获取结果 ---
        torch.cuda.synchronize() # 再次同步以确保所有评估计算已完成
        end_time = time.time()
        inference_time = end_time - start_time
        # 获取峰值显存占用（转换为MB）
        peak_memory_mb = torch.cuda.max_memory_allocated(self.device) / (1024 * 1024)
        logger.info(f"推理完成。耗时: {inference_time:.4f} 秒, 峰值显存: {peak_memory_mb:.2f} MB")

        # --- 步骤 4: 打印和保存结果 ---
        # (这部分的日志打印保持不变)
        logger.info("\n--- 新数据集上的评估结果 ---")
        if self.test_idx.numel() > 0:
            logger.info(f"测试集损失: {test_loss:.4f}")
            logger.info(f"测试集指标: {', '.join([f'{k}: {v:.4f}' for k, v in test_metrics.items()])}")
        # ... (其他 val 和 train 的日志打印) ...

        # --- 修改：合并评估指标和性能指标，然后写入CSV ---
        csv_file_path = f"CWE_test_result/{self.args.dataset}_test.csv"
        output_dir = os.path.dirname(csv_file_path)
        os.makedirs(output_dir, exist_ok=True)

        # 创建一个性能指标字典
        performance_metrics = {
            'Inference_Time_s': f"{inference_time:.4f}",
            'Peak_GPU_Memory_MB': f"{peak_memory_mb:.2f}",
            'GNN_Parameters': gnn_params,
            'Total_Parameters': total_params,
            'Total_Nodes_Evaluated': self.n_node,
        }
        
        # 将 test_metrics 和 performance_metrics 合并到一个字典中
        combined_results = test_metrics.copy()
        combined_results.update(performance_metrics)

        # 使用合并后的结果来定义CSV的头部和行数据
        fieldnames = list(combined_results.keys())
        row_data = combined_results

        file_exists = os.path.isfile(csv_file_path)

        try:
            with open(csv_file_path, mode='a', newline='') as csvfile:
                # 注意：这里需要确保 fieldnames 的顺序对于每次写入是固定的，
                # sorted(combined_results.keys()) 是一个更稳妥的选择，以防字典顺序变化
                writer = csv.DictWriter(csvfile, fieldnames=sorted(fieldnames))

                if not file_exists:
                    writer.writeheader()

                writer.writerow(row_data)
            logger.info(f"测试评估结果和性能指标已保存到: {csv_file_path}")
        except Exception as e:
            logger.error(f"保存测试结果到 CSV 时出错: {e}")
            
        return test_metrics

    def count_params(self, grad_only=True):
        params = self.get_params(grad_only = grad_only)
        return sum([p.numel() for p in params])     
    
    def to_device(self, item):
        if item != None:
            item.to(self.device)


    def run(self, n_running, rseed, prune_tolerate=1):
        evaluator_wrapper = lambda pred, labels: self.evaluator.eval(
            {"y_pred": pred.argmax(dim=-1, keepdim=True), "y_true": labels} #onehot to cls
        )["acc"]

        # kd mode
        mode = self.args.kd_mode
        # define model and optimizer
        self.init_lm()
        self.init_gnn()
        if(self.args.lm_x):
            self.init_optimizer_lm()
            self.model_lm.bert_model.gradient_checkpointing_enable()
        # training loop
        start_ep = 0
        last_is_e2e = True
        total_time = 0
        best_val_acc, final_test_acc, best_val_loss = 0, 0, float("inf")
        best_metric, final_test_metric = 0, {}
        metric_key = 'acc' if self.args.num_labels>2 else 'f1'  
        final_pred = None
 
        
        # 初始化空的DataFrame用于记录指标
        metrics_df = []
        
        if self.args.proceed:
            start_ep, last_is_e2e, best_val_acc = self.load_stat()
        
        logger.info(f"Number of all params: {self.count_params(grad_only=False)}")
        self.to_device(self.model_gnn)
        self.to_device(self.model_lm)
        
        if self.args.wu_lm > 0 and start_ep==0:
            logger.info("Warming up LM")
            self.train_lm()

        epoch = start_ep + 1
        
        while epoch < self.args.n_epochs + 1:
            is_e2e = self.args.ftmask[epoch] and not self.args.use_external_feat   
            if is_e2e: prune_tolerate -= 1
            if last_is_e2e != is_e2e:
                last_is_e2e = is_e2e
                if epoch >= self.args.peft_start > 0:
                    self.switch_to('gnn_lora')
                if is_e2e:
                    self.switch_to('e2e')
                else:
                    self.switch_to('gnn_only')
                    
            self.adjust_learning_rate(epoch)
            
            tic = time.time()   
            self.train(epoch, evaluator_wrapper, is_e2e, mode=mode)
            toc = time.time()
            total_time += toc - tic
            gc.collect()
            
            train_metrics, val_metrics, test_metrics, train_loss, val_loss, test_loss, pred = \
                self.evaluate(evaluator_wrapper,epoch, is_e2e)
                                                   
            val_metric = val_metrics[metric_key]
            
            if self.trial and prune_tolerate == 0:
                if val_metric < self.args.expected_valid_acc or self.trial.should_prune():
                    logger.critical(
                        f"valid {metric_key} {val_metric:.4f} is lower than expected {self.args.expected_valid_acc:.4f}"
                    )
                    raise optuna.exceptions.TrialPruned()

            peak_memuse = torch.cuda.max_memory_allocated(self.device) / float(1024**3)
            logger.info("Peak memuse {:.2f} G".format(peak_memuse))

            # if val_metric > best_metric:
            #     best_val_loss = val_loss
            #     best_metric = val_metric
            #     final_test_metric = test_metrics
            #     final_pred = pred
            #     if val_metric > 0.6:
            #         self.save_stat(epoch, is_e2e, val_metric, f'best_{rseed}')
            #         self.save_pred(final_pred, f'best_{rseed}')
            #         logger.info(f'best_{rseed} at ep{epoch} saved')
            # --- 早停逻辑开始 ---
            if val_metric > best_metric:
                # 效果有提升，更新所有最佳记录
                best_val_loss = val_loss
                best_metric = val_metric
                final_test_metric = test_metrics
                final_pred = pred
                
                # 因为这是目前最好的模型，所以直接保存，不再需要额外的判断
                self.save_stat(epoch, is_e2e, best_metric, f'best_{rseed}')
                self.save_pred(final_pred, f'best_{rseed}')
                logger.info(f'在第 {epoch} 轮发现新的最佳模型并已保存 (验证集-{metric_key}: {best_metric:.4f})')
            
            # 记录当前epoch的指标
            epoch_data = {
                'epoch': epoch,
                # 'is_e2e': is_e2e,
                'total_time': total_time,
                'train_loss': train_loss,
                'val_loss': val_loss,
                'test_loss': test_loss,
                # 'best_val_metric': best_metric,
                'final_test_metric': final_test_metric,
                'peak_memuse': peak_memuse
            }
            
            # 添加动态指标
            for k, v in train_metrics.items():
                epoch_data[f'train_{k}'] = v
            for k, v in val_metrics.items():
                epoch_data[f'val_{k}'] = v
            for k, v in test_metrics.items():
                epoch_data[f'test_{k}'] = v
            
            metrics_df.append(epoch_data)
            
            if epoch == self.args.n_epochs or epoch % self.args.log_every == 0:
                logger.info(
                    f"Run: {n_running}/{self.args.n_runs}/{rseed}, Epoch: {epoch}/{self.args.n_epochs}, Average epoch time: {total_time / (epoch-start_ep):.2f}"
                    f"\nTrain/Val/Test loss: {train_loss:.4f}/{val_loss:.4f}/{test_loss:.4f}"
                    f"\nTrain/Val/Test/Best val/Final test {metric_key}: {val_metric:.4f}/{test_metrics[metric_key]:.4f}/{best_metric:.4f}"
                    f"\nTrain: {', '.join([f'{k}: {v:.4f}' for k, v in train_metrics.items()])}"
                    f"\nVal: {', '.join([f'{k}: {v:.4f}' for k, v in val_metrics.items()])}"
                    f"\nTest: {', '.join([f'{k}: {v:.4f}' for k, v in test_metrics.items()])}"
                    f"\nFinal: {', '.join([f'{k}: {v:.4f}' for k, v in final_test_metric.items()])}"
                )

            epoch += 1

        # 训练结束后保存CSV文件
        csv_filename = os.path.join(
            self.args.save, 
            f"training_metrics_run{n_running}_seed{rseed}.csv"
        )
        df = pd.DataFrame(metrics_df)
        df.to_csv(csv_filename, index=False)
        logger.info(f"Training metrics saved to {csv_filename}")

        logger.info("*" * 50)
        logger.info(f"Best val {metric_key}: {best_metric}\nFinal: {', '.join([f'{k}: {v:.4f}' for k, v in final_test_metric.items()])}")
        logger.info("*" * 50)
        # self.save_pred(pred, f'last_{rseed}')

        self.args.proceed = False
        return best_metric, final_test_metric
    
def main():
    set_logging()
    gbc = LM_GNN(parse_args())
    gbc.init_lm()
    gbc.load_data()
    gbc.prepare()
    logger.info(gbc.args.save)
    # print("success")
    logger.info(gbc.args)

    if gbc.args.test_only:
        # 检查是否提供了模型路径和新数据集路径
        if not gbc.args.model_path:
            raise ValueError("在评估模式下，必须同时提供 --model_path 和 --new_data_path。")
        
        # 直接调用 evaluate_standalone 进行评估
        gbc.evaluate_standalone(gbc.args.model_path)

    else:
  
        val_accs, test_accs = [], []
        
        for i in range(gbc.args.n_runs):
            rseed = gbc.args.seed + i
            seed(rseed)
            gbc.init_loader()
            val_acc, test_acc = gbc.run(i + 1, rseed, gbc.args.prune_tolerate)
            val_accs.append(val_acc)
            #test_accs.append(test_acc)
            test_accs.append(test_acc['acc'])
        
        logger.info(f"Runned {gbc.args.n_runs} times")
        logger.info("Val Accs:")
        logger.info(val_accs)
        logger.info("Test Accs:")
        logger.info(test_accs)
        logger.info(f"Average val accuracy: {np.mean(val_accs)} ± {np.std(val_accs)}")
        logger.info(f"Average test accuracy: {np.mean(test_accs)} ± {np.std(test_accs)}")
        logger.info(f"Number of all params: {gbc.count_params(grad_only=False)}")
        logger.info(f"Number of trainable params: {gbc.count_params()}")
    # gbc.gen_model()
    # gbc.train_lm()
if __name__ == "__main__":
    main()