REDDA/TopoEncoder.py

import os
import math
import json
import argparse
import random
import pickle

from tqdm import tqdm
import numpy as np
import pandas as pd
from typing import Dict, Tuple, List, Optional
import shutup
shutup.please()
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.data import HeteroData
from torch_geometric.nn import HGTConv
from torch_geometric.loader import NeighborLoader


TARGET_NODE_TYPES = ['drug', 'protein', 'disease', 'effect', 'pathway', 'bioprocess', 'cellcomp', 'molfunc']

TYPE_ALIASES = {
    'drug': 'drug',
    'gene/protein': 'protein', 'protein': 'protein',
    'disease': 'disease',
    'side effect': 'effect', 'adverse event': 'effect', 'effect': 'effect',
    'pathway': 'pathway',
    'biological process': 'bioprocess', 'bioprocess': 'bioprocess',
    'cellular component': 'cellcomp', 'cellcomp': 'cellcomp',
    'molecular function': 'molfunc', 'molfunc': 'molfunc',
}

def norm_type(s: str) -> Optional[str]:
    if pd.isna(s): return None
    k = str(s).strip().lower()
    return TYPE_ALIASES.get(k, None)  # 未知类型直接丢弃

def set_seed(seed: int = 42):
    random.seed(seed); np.random.seed(seed)
    torch.manual_seed(seed); torch.cuda.manual_seed_all(seed)

def ensure_dir(p: str):
    if p and not os.path.exists(p): os.makedirs(p, exist_ok=True)

def sanitize_rel(rel: str) -> str:
    r = str(rel).strip().lower().replace(' ', '_')
    return r

# -----------------------
# 从“富列 CSV”构建 HeteroData
# -----------------------

def build_heterodata_from_rich_csv(kg_csv: str,
                                   dedup: bool = True) -> Tuple[HeteroData, Dict[str, Dict[str, int]]]:
    usecols = ['relation','x_id','x_type','y_id','y_type']
    df = pd.read_csv(kg_csv, usecols=usecols)

    # 归一化类型
    df['src_type'] = df['x_type'].apply(norm_type)
    df['dst_type'] = df['y_type'].apply(norm_type)
    df['rel'] = df['relation'].apply(sanitize_rel)
    # 过滤掉无法识别的类型
    keep = df['src_type'].isin(TARGET_NODE_TYPES) & df['dst_type'].isin(TARGET_NODE_TYPES)
    df = df[keep].copy().reset_index(drop=True)

    # 收集每个类型出现过的所有字符串 ID（来自 x_id 与 y_id），并映射到 0..N-1
    id2idx: Dict[str, Dict[str, int]] = {}
    for t in TARGET_NODE_TYPES:
        ids_x = df.loc[df['src_type'] == t, 'x_id']
        ids_y = df.loc[df['dst_type'] == t, 'y_id']

        # 先拼接，再去 NaN，再转字符串，最后 unique 得到 Index
        vals = pd.concat([ids_x, ids_y], ignore_index=True)
        vals = vals.dropna()
        vals = vals.astype(str)

        ids_t = pd.Index(vals.unique())  # -> Index([...])
        id2idx[t] = {sid: i for i, sid in enumerate(ids_t)}

    # 用映射把原始 x_id/y_id 转成 per-type 连续索引
    df['x_idx'] = np.nan
    df['y_idx'] = np.nan
    for t in TARGET_NODE_TYPES:
        mx = (df['src_type'] == t)
        my = (df['dst_type'] == t)
        if mx.any():
            df.loc[mx, 'x_idx'] = df.loc[mx, 'x_id'].astype(str).map(id2idx[t])
        if my.any():
            df.loc[my, 'y_idx'] = df.loc[my, 'y_id'].astype(str).map(id2idx[t])

    # 去掉映射失败的边（极少见：原表里有空 ID 或脏数据）
    df = df[df['x_idx'].notna() & df['y_idx'].notna()].copy()
    df['x_idx'] = df['x_idx'].astype(np.int64)
    df['y_idx'] = df['y_idx'].astype(np.int64)

    # 构建 HeteroData
    data = HeteroData()
    # 设置每种类型的节点数
    for t in TARGET_NODE_TYPES:
        data[t].num_nodes = len(id2idx[t])

    # 依 (src_type, rel, dst_type) 分组建边
    for (s_t, rel, d_t), g in df.groupby(['src_type','rel','dst_type']):
        if len(g) == 0: continue
        arr = g[['x_idx','y_idx']].to_numpy(dtype=np.int64)
        if dedup:
            arr = np.unique(arr, axis=0)  # 去重
        ei = torch.from_numpy(arr.T).long().contiguous()
        data[(s_t, rel, d_t)].edge_index = ei

    # 加反向边
    data = add_reverse_edges(data)
    # 度特征：log(1+deg)，用于“默认特征”
    degs = compute_hetero_degree(data)
    for t, deg in degs.items():
        data[t].deg = deg

    return data, id2idx

def log_hetero_stats(data: HeteroData, title: str = "HeteroData Stats"):
    print(f"\n=== {title} ===")
    # 节点数
    for t in data.node_types:
        print(f"Nodes[{t}]: {data[t].num_nodes}")
    # 边数（按三元组类型）
    total_edges = 0
    for (s_t, rel, d_t) in data.edge_types:
        e = int(data[(s_t, rel, d_t)].edge_index.size(1))
        print(f"Edges[{s_t} -{rel}-> {d_t}]: {e}")
        total_edges += e
    print(f"Total node types: {len(data.node_types)} | "
          f"Total edge types: {len(data.edge_types)} | "
          f"Total edges: {total_edges}\n")

def add_reverse_edges(data: HeteroData) -> HeteroData:
    edge_types = list(data.edge_types)
    for (s_t, rel, d_t) in edge_types:
        ei = data[(s_t, rel, d_t)].edge_index
        rev = (d_t, rel + "_rev", s_t)
        if rev not in data.edge_types:
            data[rev].edge_index = ei.flip(0).contiguous()
    return data

def compute_hetero_degree(data: HeteroData) -> Dict[str, torch.Tensor]:
    degs: Dict[str, torch.Tensor] = {}
    for t in data.node_types:
        n = data[t].num_nodes
        deg_total = torch.zeros(n, dtype=torch.float32)
        for (s_t, rel, d_t) in data.edge_types:
            ei = data[(s_t, rel, d_t)].edge_index
            if s_t == t and ei.numel() > 0:
                deg_total.index_add_(0, ei[0], torch.ones(ei.size(1)))
            if d_t == t and ei.numel() > 0:
                deg_total.index_add_(0, ei[1], torch.ones(ei.size(1)))
        degs[t] = torch.log1p(deg_total).unsqueeze(-1)  # [n,1]
    return degs

# -----------------------
# HGT 编码器（默认特征：类型向量 + 度投影）
# -----------------------

class HGTEncoder(nn.Module):
    def __init__(self, metadata, hidden_dim: int = 256, num_layers: int = 2, heads: int = 4, dropout: float = 0.1):
        super().__init__()
        self.metadata = metadata
        self.node_types = metadata[0]
        self.hidden_dim = hidden_dim

        # 每类节点共享的类型向量（保持归纳：新节点不会用 ID embedding）
        self.type_vec = nn.ParameterDict({t: nn.Parameter(torch.randn(1, hidden_dim) * 0.02)
                                          for t in self.node_types})
        self.deg_proj = nn.ModuleDict({t: nn.Linear(1, hidden_dim, bias=False) for t in self.node_types})

        self.convs = nn.ModuleList([
            HGTConv(hidden_dim, hidden_dim, metadata, heads=heads)
            for _ in range(num_layers)
        ])
        self.norms = nn.ModuleList([nn.ModuleDict({t: nn.LayerNorm(hidden_dim) for t in self.node_types})
                                    for _ in range(num_layers)])
        self.dropout = nn.Dropout(dropout)

    def forward(self, data: HeteroData) -> Dict[str, torch.Tensor]:
        h = {}
        for t in self.node_types:
            n = data[t].num_nodes
            tv = self.type_vec[t].expand(n, -1)
            deg = data[t].deg.to(tv.device)
            h[t] = tv + self.deg_proj[t](deg)  # 默认特征：类型向量 + 度投影

        for l, conv in enumerate(self.convs):
            h = conv(h, data.edge_index_dict)
            for t in h:
                h[t] = self.norms[l][t](h[t])
                h[t] = F.gelu(h[t])
                h[t] = self.dropout(h[t])
        return h

# -----------------------
# 关系打分（用于边重构）
# -----------------------

class DistMultRelScorer(nn.Module):
    def __init__(self, edge_types: List[Tuple[str,str,str]], hidden_dim: int):
        super().__init__()
        self.rel_vec = nn.ParameterDict()
        for (s, r, d) in edge_types:
            self.rel_vec[r] = nn.Parameter(torch.randn(hidden_dim) * 0.02)

    def score_pairs(self, h_s: torch.Tensor, h_d: torch.Tensor, rel: str) -> torch.Tensor:
        r = self.rel_vec[rel]
        return torch.sum(h_s * r * h_d, dim=-1)  # [B]

# -----------------------
# 采样器
# -----------------------

def sample_pos_neg_edges(ei: torch.Tensor, num_src: int, num_dst: int, num_samples: int) -> Tuple[torch.Tensor, torch.Tensor]:
    E = ei.size(1)
    num_samples = min(num_samples, E)
    pos_idx = torch.randint(0, E, (num_samples,), device=ei.device)
    pos_e = ei[:, pos_idx]  # [2, B]
    neg_src = torch.randint(0, num_src, (num_samples,), device=ei.device)
    neg_dst = torch.randint(0, num_dst, (num_samples,), device=ei.device)
    neg_e = torch.stack([neg_src, neg_dst], dim=0)
    return pos_e, neg_e

def make_loaders(data: HeteroData,
                 seed_types: List[str] = ('drug', 'protein'),
                 batch_size: int = 2048,
                 num_neighbors: List[int] = [10, 5],
                 shuffle: bool = True):
    loaders = {}
    for t in seed_types:
        if t not in data.node_types or data[t].num_nodes == 0:
            continue
        seeds = torch.arange(data[t].num_nodes)  # 全部该类型节点作采样起点
        loaders[t] = NeighborLoader(
            data,
            num_neighbors=num_neighbors,         # 每层邻居数，可按显存调小/大
            input_nodes=(t, seeds),
            batch_size=batch_size,
            shuffle=shuffle
        )
    return loaders

# -----------------------
# 训练（HGMAE 风格：仅边重构；如有外部特征可再加“特征重构”）
# -----------------------

def train_hgmae(data: HeteroData,
                hidden_dim: int = 256,
                layers: int = 2,
                heads: int = 4,
                edge_samples_per_rel: int = 4096,
                lr: float = 2e-3,
                weight_decay: float = 1e-4,
                epochs: int = 50,
                device: str = 'cuda' if torch.cuda.is_available() else 'cpu',
                out_dir: str = 'out_emb',
                id2idx: Optional[Dict[str, Dict[str, int]]] = None,
                save_only: Tuple[str, ...] = ('drug', 'protein')):

    ensure_dir(out_dir)
    metadata = data.metadata()

    model = HGTEncoder(metadata, hidden_dim, layers, heads).to(device)
    rel_scorer = DistMultRelScorer(metadata[1], hidden_dim).to(device)

    opt = torch.optim.AdamW(list(model.parameters()) + list(rel_scorer.parameters()),
                            lr=lr, weight_decay=weight_decay)
    bce = nn.BCEWithLogitsLoss()

    loaders = make_loaders(
        data,
        seed_types=('drug','protein'),      # 可按需扩展
        batch_size=2048,                    # 显存吃紧就降，比如 1024/512
        num_neighbors=[10, 5],              # 邻居数也可降到 [10] 或 [10,3]
        shuffle=True
    )
    seed_order = list(loaders.keys())
    assert len(seed_order) > 0, "No valid seed types to build loaders."

    for epoch in range(1, epochs + 1):
        model.train()
        running_loss, running_batches = 0.0, 0

        # 轮训不同类型的 loader（让训练更均衡）
        for seed_t in seed_order:
            loader = loaders[seed_t]
            for batch in loader:
                batch = batch.to(device, non_blocking=True)

                # 前向（仅在 batch 子图上）
                h = model(batch)

                # 边重构损失（只在 batch 的边上算；跳过反向边）
                edge_loss = torch.zeros([], device=device)
                edge_cnt = 0
                for (s_t, rel, d_t) in batch.edge_types:
                    if rel.endswith('_rev'):
                        continue
                    ei = batch[(s_t, rel, d_t)].edge_index
                    if ei.numel() == 0:
                        continue

                    # 这里的 num_nodes 使用子图的节点数
                    pos_e, neg_e = sample_pos_neg_edges(
                        ei,
                        batch[s_t].num_nodes,
                        batch[d_t].num_nodes,
                        edge_samples_per_rel
                    )

                    s_pos = rel_scorer.score_pairs(h[s_t][pos_e[0]], h[d_t][pos_e[1]], rel)
                    s_neg = rel_scorer.score_pairs(h[s_t][neg_e[0]], h[d_t][neg_e[1]], rel)
                    logit = torch.cat([s_pos, s_neg], dim=0)
                    label = torch.cat([torch.ones_like(s_pos), torch.zeros_like(s_neg)], dim=0)

                    edge_loss = edge_loss + bce(logit, label)
                    edge_cnt += 1

                if edge_cnt > 0:
                    edge_loss = edge_loss / edge_cnt

                loss = edge_loss
                opt.zero_grad(set_to_none=True)
                loss.backward()
                # torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                opt.step()

                running_loss += loss.item()
                running_batches += 1

        avg_loss = running_loss / max(1, running_batches)
        print(f"[Epoch {epoch:03d}] avg_loss={avg_loss:.4f}")

    # ===== 训练完导出 embedding（仅保存指定类型，且保存为 {orig_id: embedding} 的 dict）=====
    torch.save({'state_dict': model.state_dict(),
                'metadata': metadata,
                'hidden_dim': hidden_dim,
                'layers': layers,
                'heads': heads},
               os.path.join(out_dir, 'hgt_encoder.pt'))

    # 推理阶段：用 NeighborLoader 分批推理，按类型收集
    model.eval()
    with torch.no_grad():
        # 只对 save_only 里的类型做收集
        h_all = {t: torch.zeros((data[t].num_nodes, hidden_dim), dtype=torch.float32)
                 for t in save_only if t in data.node_types}

        # 分类型推理（不打乱）
        from torch_geometric.loader import NeighborLoader
        def make_single_loader(ntype: str, bs: int = 4096, fans: List[int] = [10,5]):
            seeds = torch.arange(data[ntype].num_nodes)
            return NeighborLoader(
                data,
                num_neighbors=fans,
                input_nodes=(ntype, seeds),
                batch_size=bs,
                shuffle=False
            )

        for ntype in list(h_all.keys()):
            loader = make_single_loader(ntype)
            for batch in loader:
                # 子图里该类型节点对应原图的索引（注意：不同版本字段名可能是 .n_id 或 .input_id）
                if hasattr(batch[ntype], 'n_id'):
                    orig_idx = batch[ntype].n_id
                elif hasattr(batch[ntype], 'input_id'):
                    orig_idx = batch[ntype].input_id
                else:
                    raise AttributeError(f"Cannot find original indices for node type '{ntype}'")

                batch = batch.to(device, non_blocking=True)
                h = model(batch)
                h_cpu = h[ntype].detach().cpu()
                h_all[ntype][orig_idx] = h_cpu  # 回填到全量矩阵

    # 需要把 per-type 连续索引 → 原始字符串ID
    # id2idx: {type: {orig_id(str): idx(int)}}
    assert id2idx is not None, "id2idx mapping is required to export {orig_id: embedding} dict."
    idx2id = {t: {idx: oid for oid, idx in mapping.items()} for t, mapping in id2idx.items() if t in h_all}

    # 组装 dict 并保存
    emb_dir = os.path.join(out_dir, 'embeddings'); os.makedirs(emb_dir, exist_ok=True)

    emb_dict_all = {}  # {'drug': {orig_id: np.ndarray}, 'protein': {...}}
    for t in h_all:
        t_vec = h_all[t].numpy()  # [N, D]
        mapping = idx2id[t]
        # 构建 {orig_id: embedding(np.ndarray)}
        emb_dict = {mapping[i]: t_vec[i] for i in range(t_vec.shape[0])}
        emb_dict_all[t] = emb_dict

        # 单独各存一份
        torch.save(emb_dict, os.path.join(emb_dir, f'{t}_emb.pt'))  # Python 端最通用
        # 如需 .npz 版本（可被 numpy 读取）：
        # np.savez_compressed(os.path.join(emb_dir, f'{t}_emb.npz'),
        #                     ids=np.array(list(emb_dict.keys()), dtype=object),
        #                     vecs=np.stack(list(emb_dict.values()), axis=0))

    # 合并存一份
    torch.save(emb_dict_all, os.path.join(emb_dir, 'topo_embeddings.pt'))
    print(f"[DONE] Saved dict embeddings for {list(h_all.keys())} to {emb_dir}")
    return model, emb_dict_all

# -----------------------
# 前向编码（用于增量并网后的新节点向量）
# -----------------------

@torch.no_grad()
def infer_embeddings_for_types(
    kg_csv: str,
    encoder_ckpt: str,
    types: Tuple[str, ...] = ('drug','protein'),
    batch_size: int = 4096,
    num_neighbors: List[int] = [10,5],
    device: str = 'cuda' if torch.cuda.is_available() else 'cpu',
    # 新增：只对这些原始ID做推理；例如 {'drug': [...], 'protein': [...]}
    target_ids: Optional[Dict[str, List[str]]] = None,
) -> Tuple[Dict[str, Dict[str, np.ndarray]], Dict[str, Dict[str,int]]]:
    """
    仅对 target_ids 指定的节点做推理（而不是全量）。
    返回：
      emb_dict_all: { 'drug': {orig_id: vec}, 'protein': {...} }（仅包含目标ID）
      id2idx:      { type: {orig_id: idx} }
    """
    # 1) 构图 + id 映射
    data, id2idx = build_heterodata_from_rich_csv(kg_csv, dedup=True)
    # 图常驻 CPU；batch 时拷到 GPU
    for et in data.edge_types: data[et].edge_index = data[et].edge_index
    for t in data.node_types:  data[t].deg = data[t].deg

    # 2) 加载编码器
    ckpt = torch.load(encoder_ckpt, map_location=device)
    model = HGTEncoder(ckpt['metadata'],
                       hidden_dim=ckpt['hidden_dim'],
                       num_layers=ckpt['layers'],
                       heads=ckpt['heads']).to(device)
    model.load_state_dict(ckpt['state_dict'])
    model.eval()

    # 3) 只收集目标ID
    emb_all: Dict[str, Dict[str, np.ndarray]] = {}
    if target_ids is None:
        target_ids = {}

    for ntype in types:
        want_ids = list(dict.fromkeys(target_ids.get(ntype, [])))  # 去重并保序
        if len(want_ids) == 0:
            continue
        if ntype not in data.node_types or data[ntype].num_nodes == 0:
            print(f"[WARN] Node type '{ntype}' not in graph or has 0 nodes; skip.")
            continue

        # 将目标 orig_id 映射为图内索引；过滤不存在的
        mapping = id2idx.get(ntype, {})
        seed_idx = []
        seed_id  = []
        for oid in want_ids:
            idx = mapping.get(oid, None)
            if idx is not None:
                seed_idx.append(idx)
                seed_id.append(oid)
        if len(seed_idx) == 0:
            print(f"[WARN] No target IDs of type '{ntype}' are present in the graph; skip.")
            continue

        seeds = torch.tensor(seed_idx, dtype=torch.long)
        from torch_geometric.loader import NeighborLoader
        loader = NeighborLoader(
            data,
            num_neighbors=num_neighbors,
            input_nodes=(ntype, seeds),
            batch_size=min(batch_size, len(seed_idx)),
            shuffle=False
        )

        res: Dict[str, np.ndarray] = {}
        for batch in loader:
            # 原图索引
            if hasattr(batch[ntype], 'n_id'):
                orig_idx = batch[ntype].n_id
            elif hasattr(batch[ntype], 'input_id'):
                orig_idx = batch[ntype].input_id
            else:
                raise AttributeError(f"Cannot locate original indices for node type {ntype}")

            # 仅前 batch_size 个是本批次的“种子”节点（NeighborLoader约定）
            seed_bs = getattr(batch[ntype], 'batch_size', None)
            if seed_bs is None:
                # 兜底（极少见）：认为全是种子
                seed_bs = orig_idx.shape[0]

            batch = batch.to(device, non_blocking=True)
            h = model(batch)
            h_seed = h[ntype][:seed_bs].detach().cpu().numpy()
            seed_orig_idx = orig_idx[:seed_bs].cpu().numpy()

            # 映射回 orig_id
            idx2id = {v:k for k,v in mapping.items()}
            for j in range(seed_bs):
                oid = idx2id.get(int(seed_orig_idx[j]), None)
                if oid is not None:
                    res[oid] = h_seed[j]
            del h
            torch.cuda.empty_cache()

        emb_all[ntype] = res

        # 报告没能推理到的目标ID（例如不在图里的）
        missing = [oid for oid in want_ids if oid not in res]
        if len(missing) > 0:
            print(f"[WARN] {ntype}: {len(missing)} target IDs not inferred (not present in graph or filtered). "
                  f"Examples: {missing[:5]}")

    return emb_all, id2idx


def filter_new_ids(emb_dict_all: Dict[str, Dict[str, np.ndarray]],
                   newly_added: Dict[str, List[str]]) -> Dict[str, Dict[str, np.ndarray]]:
    """
    newly_added: {'drug':[new_id1,...], 'protein':[new_id2,...]}
    只保留新增 ID 的 embedding。
    """
    out = {}
    for t, ids in newly_added.items():
        if t not in emb_dict_all: continue
        table = emb_dict_all[t]
        out[t] = {oid: table[oid] for oid in ids if oid in table}
    return out


from node_merge import (
    precompute_morgan_fp,
    load_candidate_protein_embedding,
    bulk_topN_similar_drugs,
    bulk_topN_similar_proteins,
)

def generate_topology_embeddings_from_csv(
    new_csv: str,
    kg_csv: str,
    encoder_ckpt: str,
    # drug 支持集（全网）
    drug_node_csv: str = "node/drug.csv",      # 需含 ['Inter_ID','SMILES']
    # protein 支持集（仅含有 ESM 的那部分）
    support_prot_ids_json: str = "protein_seq.json",   # 之前你的 support_set 的 keys（字符串ID）
    support_prot_vecs_npy: str = "support_prot_emb.npy",  # shape [N, D]，顺序与 json 的 keys 对齐
    candidate_prot_feat_dir: str = "/scratch/yg3191/PLAIN/data/ProteinFeature",  # {ID}.pkl 存 ESM
    # 选择参数
    topk_drug: int = 10,
    topk_prot: int = 10,
    min_sim_drug: float = 0.0,
    min_sim_prot: float = 0.0,
    # 推理设置
    device: str = "cuda" if torch.cuda.is_available() else "cpu",
    batch_size: int = 4096,
    num_neighbors: List[int] = [10,5],
    # 输出
    out_aug_kg_csv: Optional[str] = "kg_core_aug.csv",
    out_emb: Optional[str] = "out_emb/embeddings/topo_embeddings_NEW_only.pt"
) -> Dict[str, Dict[str, np.ndarray]]:
    """
    读取包含新分子/蛋白的 CSV，批量地为所有 query 计算与 support set 的相似度，
    一次性追加相似边到 KG，并用已训练好的 HGT 编码器为“新增 ID”导出拓扑向量。
    """
    # ---------- 1) 读新样本 ----------
    tbl = pd.read_csv(new_csv)

    def pick_col(cands):
        for c in cands:
            if c in tbl.columns: return c
            uc = c.upper()
            lc = c.lower()
            if uc in tbl.columns: return uc
            if lc in tbl.columns: return lc
        return None

    col_smiles = pick_col(["smiles","SMILES"])
    col_protid = pick_col(["protein_id","uniprot","prot","Inter_ID","id","target"])
    if col_smiles is None or col_protid is None:
        raise ValueError("new_csv 必须包含 SMILES 与 protein_id（列名可为 smiles/SMILES, protein_id/uniprot/prot 等）")

    import hashlib
    new_rows: List[Tuple[str, Optional[str], Optional[str]]] = []
    for _, r in tbl.iterrows():
        smi = None if pd.isna(r[col_smiles]) else str(r[col_smiles])
        pid = None if pd.isna(r[col_protid]) else str(r[col_protid])
        did = hashlib.md5(smi.encode()).hexdigest() if smi is not None else None  # 以 SMILES 生成新 drug_id
        if (smi is None) and (pid is None):
            continue
        new_rows.append((did, smi, pid))
    if len(new_rows) == 0:
        raise ValueError("new_csv 中没有可用的新样本行。")

    Q = len(new_rows)
    cand_smiles_list = [smi for (_, smi, _) in new_rows]
    cand_prot_ids   = [pid for (_, _, pid) in new_rows]
    print('Finish loading new rows')

    # ---------- 2) drug 支持集 + 指纹预计算（只做一次） ----------
    drug_df = pd.read_csv(drug_node_csv)
    if 'SMILES' not in drug_df.columns or 'Inter_ID' not in drug_df.columns:
        raise ValueError(f"{drug_node_csv} 需要列 ['Inter_ID','SMILES']")
    drug_df = drug_df.dropna(subset=['SMILES'])
    drug_id_smiles = dict(zip(drug_df['Inter_ID'].astype(str), drug_df['SMILES'].astype(str)))

    drug_fps_dict, _valid_smiles = precompute_morgan_fp(drug_id_smiles, radius=2, use_chirality=True)
    support_drug_ids = list(drug_fps_dict.keys())
    support_drug_fps = [drug_fps_dict[i] for i in support_drug_ids]
    # ---------- 3) protein 支持集（载入向量；N×D） ----------
    with open(support_prot_ids_json, 'r') as f:
        support_ids = list(json.load(f).keys())  # 与 npy 顺序对齐
    support_vecs = np.load(support_prot_vecs_npy).astype(np.float32)
    assert len(support_ids) == support_vecs.shape[0], "support ids 与 向量数不一致"
    
    # ---------- 4) 批量载入所有 query 的 protein 向量（去重后再回填） ----------
    uniq_pids = sorted({pid for pid in cand_prot_ids if pid is not None})
    pid2vec: Dict[str, np.ndarray] = {}
    for pid in uniq_pids:
        vec = load_candidate_protein_embedding(pid, candidate_prot_feat_dir, key='esmc_600m')
        if vec is not None:
            pid2vec[pid] = np.asarray(vec, dtype=np.float32)

    valid_pairs = [(i, pid2vec[pid]) for i, pid in enumerate(cand_prot_ids) if pid in pid2vec]
    if len(valid_pairs) > 0:
        idx_valid, vecs_valid = zip(*valid_pairs)
        cand_prot_mat = np.stack(vecs_valid, axis=0).astype(np.float32)  # [Q_valid, D]
        # ---- 4.1 所有 query 一次性与 support 做 cosine 相似度 Top-K ----
        prot_topk_valid = bulk_topN_similar_proteins(
            cand_prot_mat, support_ids, support_vecs, topN=topk_prot
        )  # List[len(idx_valid)]
        # 回填到完整 Q 长度列表，并做阈值过滤
        prot_topk_full: List[List[Tuple[str, float]]] = [[] for _ in range(Q)]
        for local_i, qidx in enumerate(idx_valid):
            row_list = [(nid, sc) for (nid, sc) in prot_topk_valid[local_i] if sc >= min_sim_prot]
            prot_topk_full[qidx] = row_list
    else:
        prot_topk_full = [[] for _ in range(Q)]
    print('Finish computing top proteins')

    # ---------- 5) 所有 query 的 ligand 相似度一次性计算（BulkTanimoto） ----------
    drug_topk_full = bulk_topN_similar_drugs(
        cand_smiles_list, support_drug_ids, support_drug_fps, topN=topk_drug, radius=2, use_chirality=True
    )
    # 阈值过滤
    drug_topk_full = [[(nid, sc) for (nid, sc) in lst if sc >= min_sim_drug] for lst in drug_topk_full]
    print('Finish computing top drugs')
    
    # ---------- 6) 一次性追加相似边到 KG（避免 per-row concat） ----------
    kg_df = pd.read_csv(kg_csv)
    required_cols = ['relation','display_relation','x_id','x_type','y_id','y_type']
    for c in required_cols:
        if c not in kg_df.columns:
            kg_df[c] = np.nan
    if 'weight' not in kg_df.columns:
        kg_df['weight'] = np.nan

    add_rows = []
    added_ids = {'drug': [], 'protein': []}

    for i in range(Q):
        new_drug_id, _, new_prot_id = new_rows[i]
        top_drugs = drug_topk_full[i]
        top_prots = prot_topk_full[i]

        if new_drug_id is not None and len(top_drugs) > 0:
            for nb_id, w in top_drugs:
                add_rows.append({
                    'relation': 'drug_drug',
                    'display_relation': 'similar_to',
                    'x_id': new_drug_id,
                    'x_type': 'drug',
                    'y_id': nb_id,
                    'y_type': 'drug',
                    'weight': float(w)
                })
            added_ids['drug'].append(new_drug_id)

        if new_prot_id is not None and len(top_prots) > 0:
            for nb_id, w in top_prots:
                add_rows.append({
                    'relation': 'protein_protein',
                    'display_relation': 'similar_to',
                    'x_id': new_prot_id,
                    'x_type': 'protein',
                    'y_id': nb_id,
                    'y_type': 'protein',
                    'weight': float(w)
                })
            added_ids['protein'].append(new_prot_id)

    if len(add_rows) > 0:
        add_df = pd.DataFrame(add_rows, columns=required_cols + ['weight'])
        kg_aug = pd.concat([kg_df, add_df], ignore_index=True)
        kg_aug.drop_duplicates(subset=['relation','x_id','y_id'], inplace=True, keep='first')
        kg_aug.reset_index(drop=True, inplace=True)
    else:
        kg_aug = kg_df.copy()

    if out_aug_kg_csv is not None:
        kg_aug.to_csv(out_aug_kg_csv, index=False)
    
    print('Finish adding new nodes to KG')
    # ---------- 7) 用已训练好的编码器推理并仅导出“新增 ID” ----------
    target_ids = {
        'drug':    [did for (did, smi, _) in new_rows if did is not None],
        'protein': [pid for (_, _, pid) in new_rows if pid is not None],
    }

    emb_new, _ = infer_embeddings_for_types(
        kg_csv=out_aug_kg_csv if out_aug_kg_csv is not None else kg_csv,
        encoder_ckpt=encoder_ckpt,
        types=('drug','protein'),
        batch_size=batch_size,
        num_neighbors=num_neighbors,
        device=device,
        target_ids=target_ids,   # ★ 只推理这些目标ID
    )

    if out_emb is not None:
        os.makedirs(os.path.dirname(out_emb), exist_ok=True)
        # save to pickle
        import pickle
        with open(out_emb, 'wb') as f:
            pickle.dump(emb_new, f)

    print("导出 ID 数量：", {k: len(v) for k,v in emb_new.items()})
    return emb_new


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--kg_csv', type=str, required=True, help='CSV with columns relation,x_id,x_type,y_id,y_type (plus others ignored)')
    ap.add_argument('--out_dir', type=str, default='out_emb')
    ap.add_argument('--hidden_dim', type=int, default=256)
    ap.add_argument('--layers', type=int, default=2)
    ap.add_argument('--heads', type=int, default=4)
    ap.add_argument('--edge_samples_per_rel', type=int, default=4096)
    ap.add_argument('--lr', type=float, default=2e-3)
    ap.add_argument('--weight_decay', type=float, default=1e-4)
    ap.add_argument('--epochs', type=int, default=50)
    ap.add_argument('--seed', type=int, default=42)
    args = ap.parse_args()

    set_seed(args.seed)
    data, id2idx = build_heterodata_from_rich_csv(args.kg_csv, dedup=True)
    log_hetero_stats(data, title="Built KG (with reverse edges)")
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    train_hgmae(data,
            hidden_dim=args.hidden_dim,
            layers=args.layers,
            heads=args.heads,
            edge_samples_per_rel=args.edge_samples_per_rel,
            lr=args.lr,
            weight_decay=args.weight_decay,
            epochs=args.epochs,
            device=device,
            out_dir=args.out_dir,
            id2idx=id2idx,
            save_only=('drug', 'protein'))

if __name__ == '__main__':
    main()