src/model.py

import csv
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from typing import Dict, List, Tuple
from torch.nn.utils.parametrizations import weight_norm
from torch.nn import TransformerEncoder, TransformerEncoderLayer

import esm

import pandas as pd
from tqdm import tqdm
from typing import Dict, List, Tuple

import tempfile
from pathlib import Path
import mdtraj as md

# import io
# import gzip
import os

from egnn_pytorch import EGNN

from transformers import AutoTokenizer, EsmForProteinFolding

import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# from re import search as re_search
import re


def determine_tcr_seq_vj(cdr3,V,J,chain,guess01=False):
    
    def file2dict(filename,key_fields,store_fields,delimiter='\t'):
        """Read file to a dictionary.
        key_fields: fields to be used as keys
        store_fields: fields to be saved as a list
        delimiter: delimiter used in the given file."""
        dictionary={}
        with open(filename, newline='') as csvfile:
            reader = csv.DictReader(csvfile,delimiter=delimiter)
            for row in reader:
                keys = [row[k] for k in key_fields]
                store= [row[s] for s in store_fields]

                sub_dict = dictionary
                for key in keys[:-1]:
                    if key not in sub_dict:
                        sub_dict[key] = {}
                    sub_dict = sub_dict[key]
                key = keys[-1]
                if key not in sub_dict:
                    sub_dict[key] = []
                sub_dict[key].append(store)
        return dictionary
    
    def get_protseqs_ntseqs(chain='B'):
        """returns sequence dictioaries for genes: protseqsV, protseqsJ, nucseqsV, nucseqsJ"""
        seq_dicts=[]
        for gene,type in zip(['v','j','v','j'],['aa','aa','nt','nt']):
            name = 'library/'+'tr'+chain.lower()+gene+'s_'+type+'.tsv'
            sdict = file2dict(name,key_fields=['Allele'],store_fields=[type+'_seq'])
            for g in sdict:
                sdict[g]=sdict[g][0][0]
            seq_dicts.append(sdict)
        return seq_dicts
    
    protVb,protJb,_,_ = get_protseqs_ntseqs(chain='B')
    protVa,protJa,_,_ = get_protseqs_ntseqs(chain='A')
    
    def splice_v_cdr3_j(pv: str, pj: str, cdr3: str) -> str:
        """
        pv: V gene protein sequence
        pj: J gene protein sequence
        cdr3: C-starting, F/W-ending CDR3 sequence (protein)
        Returns: The spliced full sequence (V[:lastC] + CDR3 + J suffix)
        """
        pv = (pv or "").strip().upper()
        pj = (pj or "").strip().upper()
        cdr3 = (cdr3 or "").strip().upper()

        # 1) V segment: Take the last 'C' (including the conserved C in V region)
        cpos = pv.rfind('C')
        if cpos == -1:
            raise ValueError("V sequence has no 'C' to anchor CDR3 start.")
        v_prefix = pv[:cpos]  # up to and including C

        # 2) Align CDR3's "end overlap" in J
        #    Start from the full length of cdr3, gradually shorten it, and find the longest suffix that can match in J
        j_suffix = pj  # fallback (in extreme cases)
        for k in range(len(cdr3), 0, -1):
            tail = cdr3[-k:]                 # CDR3's suffix
            m = re.search(re.escape(tail), pj)
            if m:
                j_suffix = pj[m.end():]      # Take the suffix from the matching segment
                break

        return v_prefix + cdr3 + j_suffix
        
    tcr_list = []
    for i in range(len(cdr3)):
        cdr3_ = cdr3[i]
        V_ = V[i]
        J_ = J[i]
        if chain=='A':
            protseqsV = protVa
            protseqsJ = protJa
        else:
            protseqsV = protVb
            protseqsJ = protJb
        if guess01:
            if '*' not in V_:
                V_+='*01'
            if '*' not in J_:
                J_+='*01'
        pv = protseqsV[V_]
        pj = protseqsJ[J_]
        # t = pv[:pv.rfind('C')]+  cdr3_ + pj[re_search(r'[FW]G.[GV]',pj).start()+1:]
        t = splice_v_cdr3_j(pv, pj, cdr3_)
        tcr_list.append(t)
    return tcr_list

# def negative_sampling_phla(df, neg_ratio=5, label_col='label', neg_label=0, random_state=42):
#     """
#     Create negative samples by shuffling the TCR sequences while keeping the peptide-HLA pairs intact.
#     Ensures that the generated (TCR, peptide, HLA) triplets do not exist in the original dataset.
#     """
#     negative_samples = []

#     # 正样本 triplet 集合
#     pos_triplets = set(zip(
#         df['tcra'], df['tcrb'], df['peptide'], df['HLA_full']
#     ))

#     for i in range(neg_ratio):
#         shuffled_df = df.copy()

#         tcr_cols = ['tcra', 'cdr3a_start', 'cdr3a_end', 'tcrb', 'cdr3b_start', 'cdr3b_end']
#         shuffled_tcr = df[tcr_cols].sample(frac=1, random_state=random_state + i).reset_index(drop=True)

#         for col in tcr_cols:
#             shuffled_df[col] = shuffled_tcr[col]

#         # 剔除：1) TCR 未改变的行  2) triplet 与正样本重复
#         mask_keep = []
#         for idx, row in shuffled_df.iterrows():
#             triplet = (row['tcra'], row['tcrb'], row['peptide'], row['HLA_full'])
#             if triplet in pos_triplets:
#                 mask_keep.append(False)
#             else:
#                 mask_keep.append(True)

#         shuffled_df = shuffled_df[mask_keep]
#         shuffled_df[label_col] = neg_label

#         negative_samples.append(shuffled_df)

#     negative_samples = pd.concat(negative_samples, ignore_index=True).drop_duplicates()
#     return negative_samples

import numpy as np
import pandas as pd

# def balanced_negative_sampling_phla(df, label_col='label', neg_label=0, random_state=42):
#     """
#     为每个 (peptide, HLA_full) 平衡采样负样本：
#     - 找出正样本最多的 peptide
#     - 该 peptide 的负样本数量 = 1:1，从其他 peptide 的 TCR 中采样（保持 peptide–HLA 配对）
#     - 其他 peptide 采样负样本，使每个 peptide 拥有相同总样本数
#     - 保证 peptide 与 HLA_full 始终保持配对关系
#     """
#     np.random.seed(random_state)

#     pos_df = df[df[label_col] != neg_label].copy()
#     pos_counts = pos_df['peptide'].value_counts()
#     max_peptide = pos_counts.idxmax()
#     max_pos = pos_counts.max()
#     total_target = max_pos * 2  # 每个 peptide 的最终样本数（正+负）

#     neg_samples = []

#     # 针对 max_peptide：负样本 = 1:1
#     df_other_tcrs = pos_df[pos_df['peptide'] != max_peptide][['tcra', 'tcrb', 'cdr3a_start', 'cdr3a_end', 'cdr3b_start', 'cdr3b_end']].copy()
#     neg_max = pos_df[pos_df['peptide'] == max_peptide].copy()
#     sampled_tcrs = df_other_tcrs.sample(
#         n=max_pos,
#         replace=True if len(df_other_tcrs) < max_pos else False,
#         random_state=random_state
#     ).reset_index(drop=True)
#     neg_max.update(sampled_tcrs)
#     neg_max[label_col] = neg_label
#     neg_samples.append(neg_max)

#     # 针对其他 peptides
#     for pep, n_pos in pos_counts.items():
#         if pep == max_peptide:
#             continue
#         n_neg = max(0, total_target - n_pos)
#         df_other_tcrs = pos_df[pos_df['peptide'] != pep][['tcra', 'tcrb', 'cdr3a_start', 'cdr3a_end', 'cdr3b_start', 'cdr3b_end']].copy()
#         neg_pep = pos_df[pos_df['peptide'] == pep].copy()
#         sampled_tcrs = df_other_tcrs.sample(
#             n=min(len(df_other_tcrs), n_neg),
#             replace=True if len(df_other_tcrs) < n_neg else False,
#             random_state=random_state
#         ).reset_index(drop=True)
#         sampled_tcrs = sampled_tcrs.iloc[:len(neg_pep)].copy() if len(sampled_tcrs) > len(neg_pep) else sampled_tcrs
#         neg_pep = pd.concat(
#             [neg_pep]*int(np.ceil(n_neg / len(neg_pep))), ignore_index=True
#         ).iloc[:n_neg]
#         neg_pep.update(sampled_tcrs)
#         neg_pep[label_col] = neg_label
#         neg_samples.append(neg_pep)

#     neg_df = pd.concat(neg_samples, ignore_index=True)
#     final_df = pd.concat([pos_df, neg_df], ignore_index=True).reset_index(drop=True)

#     return final_df

def negative_sampling_phla(df, neg_ratio=5, label_col='label', neg_label=0, random_state=42):
    """
    Create negative samples by shuffling TCRs while keeping peptide–HLA pairs intact.
    Ensures negative samples count = neg_ratio × positive samples count.
    """
    np.random.seed(random_state)
    pos_triplets = set(zip(df['tcra'], df['tcrb'], df['peptide'], df['HLA_full']))
    tcr_cols = ['tcra', 'cdr3a_start', 'cdr3a_end', 'tcrb', 'cdr3b_start', 'cdr3b_end']

    n_pos = len(df)
    target_n_neg = n_pos * neg_ratio
    all_neg = []
    
    i = 0
    while len(all_neg) < target_n_neg:
        shuffled_df = df.copy()
        shuffled_tcr = df[tcr_cols].sample(frac=1, random_state=random_state + i).reset_index(drop=True)
        for col in tcr_cols:
            shuffled_df[col] = shuffled_tcr[col]

        mask_keep = []
        for idx, row in shuffled_df.iterrows():
            triplet = (row['tcra'], row['tcrb'], row['peptide'], row['HLA_full'])
            mask_keep.append(triplet not in pos_triplets)
        shuffled_df = shuffled_df[mask_keep]
        shuffled_df[label_col] = neg_label

        all_neg.append(shuffled_df)
        i += 1

        if len(pd.concat(all_neg)) > target_n_neg * 1.5:
            break

    negative_samples = pd.concat(all_neg, ignore_index=True).drop_duplicates()
    negative_samples = negative_samples.sample(
        n=min(len(negative_samples), target_n_neg), random_state=random_state
    ).reset_index(drop=True)

    return negative_samples

# def negative_sampling_tcr(df, neg_ratio=5, label_col='label', neg_label=0, random_state=42):
#     """
#     Create negative samples by keeping TCR fixed but assigning random (peptide, HLA_full)
#     pairs that do not exist in the original dataset.
#     Ensures that the generated (TCR, peptide, HLA) triplets do not exist in the original data.
#     """
#     np.random.seed(random_state)
#     negative_samples = []

#     pos_triplets = set(zip(df['tcra'], df['tcrb'], df['peptide'], df['HLA_full']))

#     all_pairs = list(set(zip(df['peptide'], df['HLA_full'])))

#     for i in range(neg_ratio):
#         neg_df = df.copy()

#         # 随机打乱 peptide–HLA 对，但保证不会选原来的那一个
#         new_pairs = []
#         for _, row in df.iterrows():
#             while True:
#                 pep, hla = all_pairs[np.random.randint(len(all_pairs))]
#                 triplet = (row['tcra'], row['tcrb'], pep, hla)
#                 if triplet not in pos_triplets:
#                     new_pairs.append((pep, hla))
#                     break

#         neg_df[['peptide', 'HLA_full']] = pd.DataFrame(new_pairs, index=neg_df.index)
#         neg_df[label_col] = neg_label
#         negative_samples.append(neg_df)

#     negative_samples = pd.concat(negative_samples, ignore_index=True).drop_duplicates()
#     return negative_samples

class EarlyStopping:
    def __init__(self, patience=10, verbose=True, delta=0.0, save_path='checkpoint.pt'):
        """
        Early stopping based on both val_loss and val_auc.
        The model is saved whenever EITHER:
            - val_loss decreases by more than delta, OR
            - val_auc increases by more than delta.
        """
        self.patience = patience
        self.verbose = verbose
        self.counter = 0
        self.early_stop = False
        self.delta = delta
        self.save_path = save_path
        
        self.best_loss = np.inf
        self.best_auc = -np.inf

    def __call__(self, val_auc, model):
        improved = False
        
        # Check auc improvement
        if val_auc > self.best_auc + self.delta:
            self.best_auc = val_auc
            improved = True

        if improved:
            self.save_checkpoint(model, val_auc)
            self.counter = 0
        else:
            self.counter += 1
            if self.verbose:
                print(f"EarlyStopping counter: {self.counter} out of {self.patience}")
            if self.counter >= self.patience:
                self.early_stop = True

    def save_checkpoint(self, model, val_auc):
        """Save current best model."""
        if self.verbose:
            print(f"Validation improved → Saving model (Score={val_auc:.4f}) to {self.save_path}")
        torch.save(model.state_dict(), self.save_path)

# ============================================================================
# ESM2 Embedding via HuggingFace
# ============================================================================
class ESM2Encoder(nn.Module):
    def __init__(self, 
                 device="cuda:0", 
                 layer=33,
                 cache_dir='/data/cache'):
        """
        Initialize an ESM2 encoder.

        Args:
            model_name (str): Name of the pretrained ESM2 model (e.g., 'esm2_t33_650M_UR50D').
            device (str): Device to run on, e.g. 'cuda:0', 'cuda:1', or 'cpu'.
            layer (int): Layer number from which to extract representations.
        """
        super().__init__()
        self.device = device
        self.layer = layer
        
        if cache_dir is None:
            cache_dir = os.path.dirname(os.path.abspath(__file__))
        self.cache_dir = cache_dir
        os.makedirs(self.cache_dir, exist_ok=True)
        
        self.model, self.alphabet = esm.pretrained.esm2_t33_650M_UR50D()
        self.batch_converter = self.alphabet.get_batch_converter()
        self.model = self.model.eval().to(device)

    def _cache_path(self, prefix):
        base_dir = os.path.dirname(os.path.abspath(__file__))
        base_dir = base_dir + "/" + self.cache_dir
        os.makedirs(base_dir, exist_ok=True)
        return os.path.join(base_dir, f"{prefix}_esm2_layer{self.layer}.pt")

    def save_obj(self, obj, path):
        """Save object to a file (no compression)."""
        torch.save(obj, path)

    def load_obj(self, path):
        """Load object from a file (no compression)."""
        return torch.load(path, map_location="cpu", weights_only=False)
    
    @torch.no_grad()
    def _embed_batch(self, batch_data):
        batch_labels, batch_strs, batch_tokens = self.batch_converter(batch_data)
        batch_tokens = batch_tokens.to(self.device)
        results = self.model(batch_tokens, repr_layers=[self.layer], return_contacts=False)
        token_representations = results["representations"][self.layer]
        batch_lens = (batch_tokens != self.alphabet.padding_idx).sum(1)
        seq_reprs = []
        for i, tokens_len in enumerate(batch_lens):
            seq_repr = token_representations[i, 1:tokens_len-1].cpu()
            seq_reprs.append(seq_repr)
        return seq_reprs

    @torch.no_grad()
    def forward(self, df, seq_col, prefix, batch_size=64, re_embed=False, cache_save=True):
        """
        Add or update embeddings for sequences in a DataFrame.
        - If there are new sequences, automatically update the dictionary and save.
        - If re_embed=True, force re-computation of all sequences.
        """
        cache_path = self._cache_path(prefix)
        emb_dict = {}

        if os.path.exists(cache_path) and not re_embed:
            print(f"[ESM2] Loading cached embeddings from {cache_path}")
            emb_dict = self.load_obj(cache_path)
        else:
            if re_embed:
                print(f"[ESM2] Re-embedding all sequences for {prefix}")
            else:
                print(f"[ESM2] No existing cache for {prefix}, will create new.")

        seqs = [str(s).strip().upper() for s in df[seq_col].tolist() if isinstance(s, str)]
        unique_seqs = sorted(set(seqs))
        new_seqs = [s for s in unique_seqs if s not in emb_dict]

        if new_seqs:
            print(f"[ESM2] Found {len(new_seqs)} new sequences → computing embeddings...")
            data = [(str(i), s) for i, s in enumerate(new_seqs)]
            for i in tqdm(range(0, len(data), batch_size), desc=f"ESM2 update ({prefix})"):
                batch = data[i:i+batch_size]
                embs = self._embed_batch(batch)
                for (_, seq), emb in zip(batch, embs):
                    emb_dict[seq] = emb.clone()
            if cache_save:
                print(f"[ESM2] Updating cache with new sequences")
                self.save_obj(emb_dict, cache_path)
        else:
            print(f"[ESM2] No new sequences for {prefix}, using existing cache")

        return emb_dict

# ============================================================================
# ESMFold (transformers)
# ============================================================================
class ESMFoldPredictorHF(nn.Module):
    def __init__(self, 
                 model_name="facebook/esmfold_v1", 
                 cache_dir=None, 
                 device='cpu', 
                 allow_tf32=True):
        super().__init__()
        self.model_name = model_name
        self.cache_dir = cache_dir
        self.device = device
        if allow_tf32:
            torch.backends.cuda.matmul.allow_tf32 = True
            torch.backends.cudnn.allow_tf32 = True

        # tokenizer and model
        print(f"Loading ESMFold model {model_name} on {device}... {'with' if cache_dir else 'without'} cache_dir: {cache_dir}")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=cache_dir)
        self.model = EsmForProteinFolding.from_pretrained(
            model_name, low_cpu_mem_usage=True, cache_dir=cache_dir
        ).eval().to(self.device)

    @torch.no_grad()
    def infer_pdb_str(self, seq: str) -> str:
        pdb_str = self.model.infer_pdb(seq)
        return pdb_str

    @torch.no_grad()
    def forward_raw(self, seq: str):
        inputs = self.tokenizer([seq], return_tensors="pt", add_special_tokens=False)
        inputs = {k: v.to(self.device) for k, v in inputs.items()}
        outputs = self.model(**inputs)
        return outputs  # ESMFoldOutput

MAX_ASA_TIEN = {
    "ALA": 129.0, "ARG": 274.0, "ASN": 195.0, "ASP": 193.0, "CYS": 167.0,
    "GLN": 225.0, "GLU": 223.0, "GLY": 104.0, "HIS": 224.0, "ILE": 197.0,
    "LEU": 201.0, "LYS": 236.0, "MET": 224.0, "PHE": 240.0, "PRO": 159.0,
    "SER": 155.0, "THR": 172.0, "TRP": 285.0, "TYR": 263.0, "VAL": 174.0,
}
SS8_INDEX = {"H":0,"B":1,"E":2,"G":3,"I":4,"T":5,"S":6,"C":7,"-":7}

class StructureFeatureExtractorNoDSSP(nn.Module):
    def __init__(self, device="cpu"):
        super().__init__()
        self.device = device

        self.in_dim = 6 + 8 + 1 + 1 + 1  # 17

        self.to(torch.device(self.device))

    @torch.no_grad()
    def _angles(self, traj):

        L = traj.n_residues

        sphi  = np.zeros(L, dtype=np.float32); cphi  = np.zeros(L, dtype=np.float32)
        spsi  = np.zeros(L, dtype=np.float32); cpsi  = np.zeros(L, dtype=np.float32)
        someg = np.zeros(L, dtype=np.float32); comeg = np.zeros(L, dtype=np.float32)

        # 1) phi: (C_{i-1}, N_i, CA_i, C_i) —— 当前残基 i 可用 atoms[1] (N_i) 来定位
        phi_idx, phi_vals = md.compute_phi(traj)          # phi_vals: (1, n_phi)
        if phi_vals.size > 0:
            for k, atoms in enumerate(phi_idx):
                res_i = traj.topology.atom(int(atoms[1])).residue.index  # N_i 所在残基
                if 0 <= res_i < L:
                    ang = float(phi_vals[0, k])
                    sphi[res_i] = np.sin(ang); cphi[res_i] = np.cos(ang)

        # 2) psi: (N_i, CA_i, C_i, N_{i+1}) —— 当前残基 i 可用 atoms[1] (CA_i)
        psi_idx, psi_vals = md.compute_psi(traj)
        if psi_vals.size > 0:
            for k, atoms in enumerate(psi_idx):
                res_i = traj.topology.atom(int(atoms[1])).residue.index  # CA_i
                if 0 <= res_i < L:
                    ang = float(psi_vals[0, k])
                    spsi[res_i] = np.sin(ang); cpsi[res_i] = np.cos(ang)

        # 3) omega: (CA_i, C_i, N_{i+1}, CA_{i+1}) —— 当前残基 i 可用 atoms[0] (CA_i)
        omg_idx, omg_vals = md.compute_omega(traj)
        if omg_vals.size > 0:
            for k, atoms in enumerate(omg_idx):
                res_i = traj.topology.atom(int(atoms[0])).residue.index  # CA_i
                if 0 <= res_i < L:
                    ang = float(omg_vals[0, k])
                    someg[res_i] = np.sin(ang); comeg[res_i] = np.cos(ang)

        angles_feat = np.stack([sphi, cphi, spsi, cpsi, someg, comeg], axis=-1)  # [L, 6]
        return angles_feat.astype(np.float32)

    @torch.no_grad()
    def _ss8(self, traj: md.Trajectory):
        ss = md.compute_dssp(traj, simplified=False)[0]
        L = traj.n_residues
        onehot = np.zeros((L, 8), dtype=np.float32)
        for i, ch in enumerate(ss):
            onehot[i, SS8_INDEX.get(ch, 7)] = 1.0
        return onehot

    @torch.no_grad()
    def _rsa(self, traj: md.Trajectory):
        asa = md.shrake_rupley(traj, mode="residue")[0]  # (L,)
        rsa = np.zeros_like(asa, dtype=np.float32)
        for i, res in enumerate(traj.topology.residues):
            max_asa = MAX_ASA_TIEN.get(res.name.upper(), None)
            rsa[i] = 0.0 if not max_asa else float(asa[i] / max_asa)
        return np.clip(rsa, 0.0, 1.0)[:, None]

    @torch.no_grad()
    def _contact_count(self, traj: md.Trajectory, cutoff_nm=0.8):
        L = traj.n_residues
        ca_atoms = traj.topology.select("name CA")
        if len(ca_atoms) == L:
            coors = traj.xyz[0, ca_atoms, :]  # nm
        else:
            xyz = traj.xyz[0]
            coors = []
            for res in traj.topology.residues:
                idxs = [a.index for a in res.atoms]
                coors.append(xyz[idxs, :].mean(axis=0))
            coors = np.array(coors, dtype=np.float32)
        diff = coors[:, None, :] - coors[None, :, :]
        dist = np.sqrt((diff**2).sum(-1))  # nm
        mask = (dist < cutoff_nm).astype(np.float32)
        np.fill_diagonal(mask, 0.0)
        cnt = mask.sum(axis=1)
        return cnt[:, None].astype(np.float32)

    @torch.no_grad()
    def _plddt(self, pdb_file: str):
        # 用 Biopython 读取 PDB 的 B-factor（ESMFold/AlphaFold 会把 pLDDT 写在这里）
        from Bio.PDB import PDBParser
        import numpy as np

        parser = PDBParser(QUIET=True)
        structure = parser.get_structure("prot", pdb_file)
        model = structure[0]

        res_plddt = []
        for chain in model:
            for residue in chain:
                atoms = list(residue.get_atoms())
                if len(atoms) == 0:
                    res_plddt.append(0.0)
                    continue
                # 该残基原子 B-factor 的均值
                bvals = [float(atom.get_bfactor()) for atom in atoms]
                res_plddt.append(float(np.mean(bvals)))

        # 归一化到 [0,1]
        plddt = np.array(res_plddt, dtype=np.float32) / 100.0
        plddt = np.clip(plddt, 0.0, 1.0)
        return plddt[:, None]  # [L,1]

    @torch.no_grad()
    def _parse_and_features(self, pdb_file: str):
        traj = md.load(pdb_file)
        L = traj.n_residues

        angles = self._angles(traj)              # [L,6]
        ss8    = self._ss8(traj)                 # [L,8]
        rsa    = self._rsa(traj)                 # [L,1]
        cnt    = self._contact_count(traj)       # [L,1]
        plddt  = self._plddt(pdb_file)           # [L,1]

        feats = np.concatenate([angles, ss8, rsa, cnt, plddt], axis=1).astype(np.float32)  # [L,17]

        ca_atoms = traj.topology.select("name CA")
        if len(ca_atoms) == L:
            coors_nm = traj.xyz[0, ca_atoms, :]
        else:
            xyz = traj.xyz[0]
            res_coords = []
            for res in traj.topology.residues:
                idxs = [a.index for a in res.atoms]
                res_coords.append(xyz[idxs, :].mean(axis=0))
            coors_nm = np.array(res_coords, dtype=np.float32)
        coors_ang = coors_nm * 10.0  # nm -> Å
        return coors_ang.astype(np.float32), feats  # [L,3], [L,17]

    @torch.no_grad()
    def forward(self, pdb_file: str):
        coors_ang, scalars = self._parse_and_features(pdb_file)
        coors = torch.tensor(coors_ang, dtype=torch.float32, device=self.device)   # [N,3]
        scalars = torch.tensor(scalars,  dtype=torch.float32, device=self.device)    # [N,17]

        return scalars, coors  # [N,17], [N,3]

class ResiduePipelineWithHFESM:
    def __init__(self, 
                 esm_model_name="facebook/esmfold_v1",
                 cache_dir=None,
                 esm_device='cpu',
                 allow_tf32=True
                 ):
        self.esm = ESMFoldPredictorHF(esm_model_name, cache_dir, esm_device, allow_tf32)
        self.struct_encoder = StructureFeatureExtractorNoDSSP(device=esm_device)
        self.cache_dir = cache_dir

    @torch.no_grad()
    def __call__(self, seq: str, save_pdb_path: str = None) -> torch.Tensor:
        pdb_str = self.esm.infer_pdb_str(seq)
        if save_pdb_path is None:
            tmpdir = self.cache_dir if self.cache_dir is not None else tempfile.gettempdir()
            save_pdb_path = str(Path(tmpdir) / "esmfold_pred_fold5.pdb")
        Path(save_pdb_path).write_text(pdb_str)

        struct_emb, struct_coords = self.struct_encoder(save_pdb_path)
        return struct_emb, struct_coords 

def sanitize_protein_seq(seq: str) -> str:
    if not isinstance(seq, str):
        return ""
    s = "".join(seq.split()).upper()
    allowed = set("ACDEFGHIKLMNPQRSTVWYXBZJUO")
    return "".join([c for c in s if c in allowed])

@torch.no_grad()
def batch_embed_to_dicts(
    df: pd.DataFrame,
    seq_col: str,
    pipeline,
    show_progress: bool = True,
) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], List[Tuple[str, str]]]:
    """
    Returns:
      - emb_dict:  {seq -> z(torch.Tensor[L, D])}
      - coord_dict:{seq -> coords(torch.Tensor[L, 3])}
      - failures:  [(seq, err_msg), ...]
    """

    raw_list = df[seq_col].astype(str).tolist()
    seqs = []
    for s in raw_list:
        ss = sanitize_protein_seq(s)
        if ss:
            seqs.append(ss)
    uniq_seqs = sorted(set(seqs)) 

    logger.info(f"Total rows: {len(df)}, valid seqs: {len(seqs)}, unique: {len(uniq_seqs)}")

    emb_dict: Dict[str, torch.Tensor] = {}
    coord_dict: Dict[str, torch.Tensor] = {}
    failures: List[Tuple[str, str]] = []

    iterator = tqdm(uniq_seqs, desc="ESMfold Predicting structure...") if show_progress else uniq_seqs
    for seq in tqdm(iterator):
        if seq in emb_dict:
            continue
        try:
            z_t, c_t = pipeline(seq)      # z: [L, D], coords: [L, 3] (torch.Tensor)
            emb_dict[seq] = z_t.detach().float().cpu()
            coord_dict[seq] = c_t.detach().float().cpu()
        except Exception as e:
            failures.append((seq, repr(e)))
            continue

    logger.info(f"[DONE] OK: {len(emb_dict)}, Failed: {len(failures)}")
    if failures[:3]:
        logger.error("[SAMPLE failures]", failures[:3])
    return emb_dict, coord_dict, failures

class ESMFoldEncoder(nn.Module):
    def __init__(self, model_name="facebook/esmfold_v1", esm_cache_dir="/data/esm_cache", cache_dir="/data/cache"):
        super(ESMFoldEncoder, self).__init__()
        self.model_name = model_name
        self.esm_cache_dir = esm_cache_dir
        self.cache_dir = cache_dir
    
    def save_obj(self, obj, path):
        """Save object to a file (no compression)."""
        torch.save(obj, path)
        
    def load_obj(self, path):
        """Load object from a file (no compression)."""
        return torch.load(path, map_location='cpu', weights_only=False)

    def load_esm_dict(self, device, df_data, chain, re_embed):

        def _clean_unique(series: pd.Series) -> list:
            cleaned = []
            for s in series.astype(str).tolist():
                ss = sanitize_protein_seq(s)
                if ss:
                    cleaned.append(ss)
            return sorted(set(cleaned))
        
        def _retry_embed_df(
            df: pd.DataFrame,
            chain: str,
            max_retries: int = 2,
            show_progress: bool = True,
        ):
            """
            Try to embed protein sequences with retries on failures.

            Args:
                df (pd.DataFrame): A DataFrame containing a column `chain` with sequences.
                chain (str): The column name containing the sequences (e.g., "alpha", "beta").
                pipeline: An embedding pipeline, should return (embedding, coords) for a sequence.
                max_retries (int): Maximum number of retries for failed sequences.
                show_progress (bool): Whether to display tqdm progress bars.

            Returns:
                feat_dict (Dict[str, torch.Tensor]): {sequence -> embedding tensor [L, D]}.
                coord_dict (Dict[str, torch.Tensor]): {sequence -> coordinate tensor [L, 3]}.
                failures (List[Tuple[str, str]]): List of (sequence, error_message) that still failed after retries.
            """
            
            pipeline = ResiduePipelineWithHFESM(
                esm_model_name=self.model_name,
                cache_dir=self.esm_cache_dir,
                esm_device=device
            )
        
            # 1. First attempt
            feat_dict, coord_dict, failures = batch_embed_to_dicts(
                df, chain, pipeline, show_progress=show_progress
            )

            # 2. Retry loop for failed sequences
            tries = 0
            while failures and tries < max_retries:
                tries += 1
                retry_seqs = [s for s, _ in failures]
                logger.info(f"[retry {tries}/{max_retries}] {len(retry_seqs)} sequences")
                retry_df = pd.DataFrame({chain: retry_seqs})

                f2, c2, failures = batch_embed_to_dicts(
                    retry_df, chain, pipeline, show_progress=show_progress
                )
                feat_dict.update(f2)
                coord_dict.update(c2)

            return feat_dict, coord_dict, failures

        def update_with_new_seqs(feat_dict, coord_dict, chain):
            base_dir = os.path.dirname(os.path.abspath(__file__))
            base_dir = base_dir + "/" + self.cache_dir
            os.makedirs(base_dir, exist_ok=True)
            path_feat = os.path.join(base_dir, f"{chain}_feat_dict.pt")
            path_coords = os.path.join(base_dir, f"{chain}_coord_dict.pt")

            all_seqs_clean = _clean_unique(df_data[chain])
            new_seqs = [s for s in all_seqs_clean if s not in feat_dict]
            if not new_seqs:
                logger.info(f"No new {chain} sequences found")
                return feat_dict, coord_dict

            logger.info(f"Found new {chain} sequences, embedding...")
            df_new = pd.DataFrame({chain: new_seqs})
            new_feat_dict, new_coord_dict, failures = _retry_embed_df(df_new, chain, max_retries=100)
            feat_dict.update(new_feat_dict)
            coord_dict.update(new_coord_dict)
            self.save_obj(feat_dict, path_feat)
            self.save_obj(coord_dict, path_coords)

            if failures:
                for seq, err in failures:
                    logger.error(f"[create] failed: {seq} | {err}")
            
            logger.info(f"Updated and saved {path_feat} and {path_coords}")
            
            return feat_dict, coord_dict

        def get_or_create_dict(chain):
            base_dir = os.path.dirname(os.path.abspath(__file__)) + "/" + self.cache_dir
            os.makedirs(base_dir, exist_ok=True)
            path_feat = os.path.join(base_dir, f"{chain}_feat_dict.pt")
            path_coords = os.path.join(base_dir, f"{chain}_coord_dict.pt")
            
            if os.path.exists(path_feat) and not re_embed:
                logger.info(f"Loading {path_feat} and {path_coords}")
                feat_dict = self.load_obj(path_feat)
                coord_dict = self.load_obj(path_coords)
            else:
                logger.info(f"{path_feat} and {path_coords} not found or re_embed=True, generating...")
                unique_seqs = _clean_unique(df_data[chain])
                df_uniq = pd.DataFrame({chain: unique_seqs})
                feat_dict, coord_dict, failures = _retry_embed_df(
                    df_uniq, chain, show_progress=True, max_retries=100
                )
                self.save_obj(feat_dict, path_feat)
                self.save_obj(coord_dict, path_coords)

                if failures:
                    for seq, err in failures:
                        logger.error(f"[create] failed: {seq} | {err}")

                logger.info(f"Saved {path_feat} and {path_coords}")
            
            return feat_dict, coord_dict

        self.dict[chain+'_feat'], self.dict[chain+'_coord'] = update_with_new_seqs(*get_or_create_dict(chain), chain)

    def pad_and_stack(self, batch_feats, L_max, batch_coors):
        """
        batch_feats: list of [L_i, D] tensors
        batch_coors: list of [L_i, 3] tensors
        return:
        feats: [B, L_max, D]
        coors: [B, L_max, 3]
        mask : [B, L_max]  (True for real tokens)
        """
        assert len(batch_feats) == len(batch_coors)
        B = len(batch_feats)
        D = batch_feats[0].shape[-1]

        feats_pad = []
        coors_pad = []
        masks = []

        for x, c in zip(batch_feats, batch_coors):
            L = x.shape[0]
            pad_L = L_max - L
            # pad feats/coors with zeros
            feats_pad.append(torch.nn.functional.pad(x, (0, 0, 0, pad_L)))       # [L_max, D]
            coors_pad.append(torch.nn.functional.pad(c, (0, 0, 0, pad_L)))       # [L_max, 3]
            m = torch.zeros(L_max, dtype=torch.bool)
            m[:L] = True
            masks.append(m)

        feats = torch.stack(feats_pad, dim=0)   # [B, L_max, D]
        coors = torch.stack(coors_pad, dim=0)   # [B, L_max, 3]
        mask  = torch.stack(masks, dim=0)       # [B, L_max]
        return feats, coors, mask
    
    def forward(self, df_data, chain, device='cpu', re_embed=False):
        """
        df_data: pd.DataFrame with a column `chain` containing sequences
        chain: str, e.g. "alpha" or "beta"
        device: str, e.g. 'cpu' or 'cuda:0'
        re_embed: bool, whether to re-embed even if cached files exist
        """
        self.dict = {}
        self.load_esm_dict(device, df_data, chain, re_embed)

        batch_feats = []
        batch_coors = []
        for seq in df_data[chain].astype(str).tolist():
            ss = sanitize_protein_seq(seq)
            if ss in self.dict[chain+'_feat'] and ss in self.dict[chain+'_coord']:
                batch_feats.append(self.dict[chain+'_feat'][ss])
                batch_coors.append(self.dict[chain+'_coord'][ss])
            else:
                raise ValueError(f"Sequence not found in embedding dict: {ss}")

        # L_max = max(x.shape[0] for x in batch_feats)

        return batch_feats, batch_coors


# =================================== Dataset / Collate ===========================================
class PepHLA_Dataset(torch.utils.data.Dataset):
    def __init__(self, df, phys_dict, esm2_dict, struct_dict):
        self.df = df
        self.phys_dict = phys_dict
        self.esm2_dict = esm2_dict
        self.struct_dict = struct_dict

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

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        pep = row['peptide']
        hla = row['HLA_full']
        label = torch.tensor(row['label'], dtype=torch.float32)

        pep_phys = self.phys_dict['pep'][pep]
        pep_esm = self.esm2_dict['pep'][pep]

        hla_phys = self.phys_dict['hla'][hla]
        hla_esm = self.esm2_dict['hla'][hla]
        hla_struct, hla_coord = self.struct_dict[hla]

        return {
            'pep_phys': pep_phys,
            'pep_esm': pep_esm,
            'hla_phys': hla_phys,
            'hla_esm': hla_esm,
            'hla_struct': hla_struct,
            'hla_coord': hla_coord,
            'label': label,
            'pep_id': pep,
            'hla_id': hla,
        }
        
def peptide_hla_collate_fn(batch):
    def pad_or_crop(x, original_len, target_len):
        L, D = x.shape
        valid_len = min(original_len, target_len)
        valid_part = x[:valid_len]
        if valid_len < target_len:
            pad_len = target_len - valid_len
            padding = x.new_zeros(pad_len, D)
            return torch.cat([valid_part, padding], dim=0)
        else:
            return valid_part

    out_batch = {}

    pep_lens = [len(item['pep_id']) for item in batch]
    max_pep_len = max(pep_lens)

    for key in batch[0].keys():
        if key == 'label':
            out_batch[key] = torch.stack([item[key] for item in batch])
        elif key.startswith('pep_') and not key.endswith('_id'):
            out_batch[key] = torch.stack([pad_or_crop(item[key], len(item['pep_id']), max_pep_len) for item in batch])
        elif key.endswith('_id'):
            out_batch[key] = [item[key] for item in batch]
        else:
            out_batch[key] = torch.stack([item[key] for item in batch])
        
    def make_mask(lengths, max_len):
        masks = []
        for L in lengths:
            m = torch.zeros(max_len, dtype=torch.bool)
            m[:L] = True
            masks.append(m)
        return torch.stack(masks)

    out_batch['pep_mask'] = make_mask(pep_lens, max_pep_len)
    return out_batch

# =================================== Dataset / Collate ===========================================
class TCRPepHLA_Dataset(torch.utils.data.Dataset):
    """
    Dataset for TCRα + TCRβ + peptide + HLA binding.
    """
    def __init__(self, df, phys_dict, esm2_dict, struct_dict, pep_hla_feat_dict):
        self.df = df
        self.phys_dict = phys_dict
        self.esm2_dict = esm2_dict
        self.struct_dict = struct_dict
        self.pep_hla_feat_dict = pep_hla_feat_dict

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

    def __getitem__(self, idx):
        row = self.df.iloc[idx]
        tcra = row['tcra']
        tcrb = row['tcrb']
        pep = row['peptide']
        hla = row['HLA_full']
        label = torch.tensor(row['label'], dtype=torch.float32)

        # ---- TCRα ----
        tcra_phys = self.phys_dict['tcra'][tcra]
        tcra_esm  = self.esm2_dict['tcra'][tcra]
        tcra_struct, tcra_coord = self.struct_dict['tcra'][tcra]
        tcra_cdr3_start = torch.tensor(row['cdr3a_start'], dtype=torch.long)
        tcra_cdr3_end   = torch.tensor(row['cdr3a_end'], dtype=torch.long)

        # ---- TCRβ ----
        tcrb_phys = self.phys_dict['tcrb'][tcrb]
        tcrb_esm  = self.esm2_dict['tcrb'][tcrb]
        tcrb_struct, tcrb_coord = self.struct_dict['tcrb'][tcrb]
        tcrb_cdr3_start = torch.tensor(row['cdr3b_start'], dtype=torch.long)
        tcrb_cdr3_end   = torch.tensor(row['cdr3b_end'], dtype=torch.long)

        # ---- peptide ----
        pep_phys = self.phys_dict['pep'][pep]
        pep_esm  = self.esm2_dict['pep'][pep]
        pep_struct, pep_coord = self.struct_dict['pep'][pep]

        # ---- HLA ----
        hla_phys = self.phys_dict['hla'][hla]
        hla_esm  = self.esm2_dict['hla'][hla]
        hla_struct, hla_coord = self.struct_dict['hla'][hla]
        
        feats = self.pep_hla_feat_dict[(pep, hla)]
        pep_feat_pretrain = feats['pep_feat_pretrain']
        hla_feat_pretrain = feats['hla_feat_pretrain']

        return {
            # TCRα
            'tcra_phys': tcra_phys,
            'tcra_esm': tcra_esm,
            'tcra_struct': tcra_struct,
            'tcra_coord': tcra_coord,
            'cdr3a_start': tcra_cdr3_start,
            'cdr3a_end': tcra_cdr3_end,

            # TCRβ
            'tcrb_phys': tcrb_phys,
            'tcrb_esm': tcrb_esm,
            'tcrb_struct': tcrb_struct,
            'tcrb_coord': tcrb_coord,
            'cdr3b_start': tcrb_cdr3_start,
            'cdr3b_end': tcrb_cdr3_end,

            # peptide
            'pep_phys': pep_phys,
            'pep_esm': pep_esm,
            'pep_struct': pep_struct,
            'pep_coord': pep_coord,

            # HLA
            'hla_phys': hla_phys,
            'hla_esm': hla_esm,
            'hla_struct': hla_struct,
            'hla_coord': hla_coord,

            'tcra_id': tcra,
            'tcrb_id': tcrb,
            'pep_id': pep,
            'hla_id': hla,
            'label': label,

            'pep_feat_pretrain': pep_feat_pretrain,
            'hla_feat_pretrain': hla_feat_pretrain,
        }

# =================================== Collate Function ===========================================
def tcr_pep_hla_collate_fn(batch):
    def pad_or_crop(x, original_len, target_len):
        L, D = x.shape
        valid_len = min(original_len, target_len)
        valid_part = x[:valid_len]
        if valid_len < target_len:
            pad_len = target_len - valid_len
            padding = x.new_zeros(pad_len, D)
            return torch.cat([valid_part, padding], dim=0)
        else:
            return valid_part

    out_batch = {}

    tcra_lens = [len(item['tcra_id']) for item in batch]
    tcrb_lens = [len(item['tcrb_id']) for item in batch]
    pep_lens  = [len(item['pep_id'])  for item in batch]

    max_tcra_len = max(tcra_lens)
    max_tcrb_len = max(tcrb_lens)
    max_pep_len  = max(pep_lens)

    for key in batch[0].keys():
        if key == 'label':
            out_batch[key] = torch.stack([item[key] for item in batch])

        elif key.startswith('tcra_') and not key.endswith('_id'):
            out_batch[key] = torch.stack([pad_or_crop(item[key], len(item['tcra_id']), max_tcra_len) for item in batch])

        elif key.startswith('tcrb_') and not key.endswith('_id'):
            out_batch[key] = torch.stack([pad_or_crop(item[key], len(item['tcrb_id']), max_tcrb_len) for item in batch])

        elif key.startswith('pep_') and not key.endswith('_id'):
            out_batch[key] = torch.stack([pad_or_crop(item[key], len(item['pep_id']), max_pep_len) for item in batch])

        elif key.endswith('_id'):
            out_batch[key] = [item[key] for item in batch]

        else:
            out_batch[key] = torch.stack([item[key] for item in batch])

    def make_mask(lengths, max_len):
        masks = []
        for L in lengths:
            m = torch.zeros(max_len, dtype=torch.bool)
            m[:L] = True
            masks.append(m)
        return torch.stack(masks)

    out_batch['tcra_mask'] = make_mask(tcra_lens, max_tcra_len)
    out_batch['tcrb_mask'] = make_mask(tcrb_lens, max_tcrb_len)
    out_batch['pep_mask']  = make_mask(pep_lens,  max_pep_len)

    return out_batch

# ==================================== 小积木：投影 + 门控 =========================================
class ResidueProjector(nn.Module):
    """把不同分支的通道维度对齐到同一 D"""
    def __init__(self, in_dim, out_dim):
        super().__init__()
        self.proj = nn.Linear(in_dim, out_dim) if in_dim != out_dim else nn.Identity()
    def forward(self, x):  # x: [B,L,Di]
        return self.proj(x)

class ResidueDoubleFusion(nn.Module):
    """
    ResidueDoubleFusion:
    A residue-level two-branch fusion module that combines two modalities (x1, x2)
    using cross-attention followed by gated residual fusion and linear projection.

    Typical usage:
        - x1: physicochemical features
        - x2: ESM embeddings  (or structure features)
    """
    def __init__(self, dim, num_heads=8, dropout=0.1):
        super().__init__()
        self.dim = dim

        # Cross-attention: allows information flow between two modalities
        self.cross_attn = nn.MultiheadAttention(
            embed_dim=dim, num_heads=num_heads, dropout=dropout, batch_first=True
        )

        # Gating mechanism: adaptively weight two modalities per residue
        self.gate = nn.Sequential(
            nn.Linear(dim * 2, dim),
            nn.ReLU(),
            nn.Linear(dim, 1),
            nn.Sigmoid()
        )

        # Optional projection after fusion
        self.out_proj = nn.Linear(dim, dim)

        # Layer norms for stable training
        self.norm_x1 = nn.LayerNorm(dim)
        self.norm_x2 = nn.LayerNorm(dim)
        self.norm_out = nn.LayerNorm(dim)

    def forward(self, x1, x2):
        """
        Args:
            x1: Tensor [B, L, D] - first modality (e.g., physicochemical)
            x2: Tensor [B, L, D] - second modality (e.g., ESM embeddings)
        Returns:
            fused: Tensor [B, L, D] - fused residue-level representation
        """

        # 1) Normalize both branches
        x1_norm = self.norm_x1(x1)
        x2_norm = self.norm_x2(x2)

        # 2) Cross-attention (x1 queries, x2 keys/values)
        #    This allows x1 to attend to x2 at each residue position
        attn_out, _ = self.cross_attn(
            query=x1_norm,
            key=x2_norm,
            value=x2_norm
        )  # [B, L, D]

        # 3) Gating between original x1 and attention-enhanced x2
        gate_val = self.gate(torch.cat([x1, attn_out], dim=-1))  # [B, L, 1]
        fused = gate_val * x1 + (1 - gate_val) * attn_out

        # 4) Optional projection + normalization
        fused = self.out_proj(fused)
        fused = self.norm_out(fused)

        return fused

class ResidueTripleFusion(nn.Module):
    """
    ResidueTripleFusion:
    A hierarchical three-branch feature fusion module for residue-level representations.
    
    Step 1: Fuse physicochemical features and protein language model embeddings.
    Step 2: Fuse the intermediate representation with structure-based features.
    
    Each fusion step uses ResidueDoubleFusion (cross-attention + gating + linear projection).
    """
    def __init__(self, dim, num_heads=8, dropout=0.1):
        super().__init__()
        # Fuse physicochemical + ESM embeddings
        self.fuse_phys_esm = ResidueDoubleFusion(dim, num_heads=num_heads, dropout=dropout)
        # Fuse the fused phys+esm representation with structure embeddings
        self.fuse_f12_struct = ResidueDoubleFusion(dim, num_heads=num_heads, dropout=dropout)

    def forward(self, phys, esm, struct):
        """
        Args:
            phys:   Tensor [B, L, D], physicochemical features (e.g., AAindex-based)
            esm:    Tensor [B, L, D], protein language model embeddings (e.g., ESM2, ProtT5)
            struct: Tensor [B, L, D], structure-derived features (e.g., torsion, RSA)
        
        Returns:
            fused:  Tensor [B, L, D], final fused representation
        """
        # Step 1: Fuse physicochemical and ESM embeddings
        f12 = self.fuse_phys_esm(phys, esm)

        # Step 2: Fuse the intermediate fused representation with structure features
        fused = self.fuse_f12_struct(f12, struct)

        return fused

class BANLayer(nn.Module):
    """
    Bilinear Attention Network Layer with proper 2D masked-softmax.
    v_mask: [B, L_v]  True=valid
    q_mask: [B, L_q]  True=valid
    """
    def __init__(self, v_dim, q_dim, h_dim, h_out, act='ReLU', dropout=0.2, k=3):
        super().__init__()
        self.c = 32
        self.k = k
        self.v_dim = v_dim
        self.q_dim = q_dim
        self.h_dim = h_dim
        self.h_out = h_out

        self.v_net = FCNet([v_dim, h_dim * self.k], act=act, dropout=dropout)
        self.q_net = FCNet([q_dim, h_dim * self.k], act=act, dropout=dropout)

        if 1 < k:
            self.p_net = nn.AvgPool1d(self.k, stride=self.k)

        if h_out <= self.c:
            self.h_mat  = nn.Parameter(torch.Tensor(1, h_out, 1, h_dim * self.k).normal_())
            self.h_bias = nn.Parameter(torch.Tensor(1, h_out, 1, 1).normal_())
        else:
            self.h_net = weight_norm(nn.Linear(h_dim * self.k, h_out), dim=None)

        self.bn = nn.BatchNorm1d(h_dim)

    def attention_pooling(self, v, q, att_map):  # att_map: [B, L_v, L_q]
        logits = torch.einsum('bvk,bvq,bqk->bk', (v, att_map, q))
        if 1 < self.k:
            logits = self.p_net(logits.unsqueeze(1)).squeeze(1) * self.k
        return logits

    def _masked_softmax_2d(self, logits, v_mask, q_mask):
        """
        logits:  [B, h_out, L_v, L_q]
        v_mask:  [B, L_v]  or None
        q_mask:  [B, L_q]  or None
        return:  probs  [B, h_out, L_v, L_q]  (masked entries=0, 在有效的二维子矩阵内归一化)
        """
        B, H, Lv, Lq = logits.shape
        device = logits.device
        if v_mask is None:
            v_mask = torch.ones(B, Lv, dtype=torch.bool, device=device)
        if q_mask is None:
            q_mask = torch.ones(B, Lq, dtype=torch.bool, device=device)

        mask2d = (v_mask[:, :, None] & q_mask[:, None, :])          # [B, Lv, Lq]
        mask2d = mask2d[:, None, :, :].expand(B, H, Lv, Lq)         # [B, H, Lv, Lq]

        logits = logits.masked_fill(~mask2d, -float('inf'))

        # 在 Lv*Lq 的联合空间做 softmax
        flat = logits.view(B, H, -1)                                # [B, H, Lv*Lq]
        # 处理极端情况：某些样本可能无有效格子，避免 NaN
        flat = torch.where(torch.isinf(flat), torch.full_like(flat, -1e9), flat)
        flat = F.softmax(flat, dim=-1)
        flat = torch.nan_to_num(flat, nan=0.0)                      # 安全兜底
        probs = flat.view(B, H, Lv, Lq)

        # 把被 mask 的位置清零（数值稳定 & 便于可视化）
        probs = probs * mask2d.float()
        return probs

    def forward(self, v, q, v_mask=None, q_mask=None, softmax=True):
        """
        v: [B, L_v, Dv], q: [B, L_q, Dq]
        """
        B, L_v, _ = v.size()
        _, L_q, _ = q.size()

        v_ = self.v_net(v)   # [B, L_v, h_dim*k]
        q_ = self.q_net(q)   # [B, L_q, h_dim*k]

        if self.h_out <= self.c:
            att_maps = torch.einsum('xhyk,bvk,bqk->bhvq', (self.h_mat, v_, q_)) + self.h_bias   # [B,H,Lv,Lq]
        else:
            v_t = v_.transpose(1, 2).unsqueeze(3)                      # [B, K, Lv, 1]
            q_t = q_.transpose(1, 2).unsqueeze(2)                      # [B, K, 1, Lq]
            d_  = torch.matmul(v_t, q_t)                               # [B, K, Lv, Lq]
            att_maps = self.h_net(d_.permute(0, 2, 3, 1))              # [B, Lv, Lq, H]
            att_maps = att_maps.permute(0, 3, 1, 2)                    # [B, H, Lv, Lq]

        if softmax:
            att_maps = self._masked_softmax_2d(att_maps, v_mask, q_mask)
        else:
            # 即使不 softmax，也把无效格子清 0，避免泄漏
            if v_mask is not None:
                att_maps = att_maps.masked_fill(~v_mask[:, None, :, None], 0.0)
            if q_mask is not None:
                att_maps = att_maps.masked_fill(~q_mask[:, None, None, :], 0.0)

        # 注意：此时 v_ / q_ 仍是 [B, L, K]，与 att_maps 的 [B,H,Lv,Lq] 对齐
        logits = self.attention_pooling(v_, q_, att_maps[:, 0, :, :])
        for i in range(1, self.h_out):
            logits = logits + self.attention_pooling(v_, q_, att_maps[:, i, :, :])

        logits = self.bn(logits)
        return logits, att_maps

class FCNet(nn.Module):
    def __init__(self, dims, act='ReLU', dropout=0.2):
        super(FCNet, self).__init__()

        layers = []
        for i in range(len(dims) - 2):
            in_dim = dims[i]
            out_dim = dims[i + 1]
            if 0 < dropout:
                layers.append(nn.Dropout(dropout))
            layers.append(weight_norm(nn.Linear(in_dim, out_dim), dim=None))
            if '' != act:
                layers.append(getattr(nn, act)())
        if 0 < dropout:
            layers.append(nn.Dropout(dropout))
        layers.append(weight_norm(nn.Linear(dims[-2], dims[-1]), dim=None))
        if '' != act:
            layers.append(getattr(nn, act)())

        self.main = nn.Sequential(*layers)

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

class StackedEGNN(nn.Module):
    def __init__(self, dim, layers, update_coors=False, **egnn_kwargs):
        super().__init__()
        self.layers = nn.ModuleList([
            EGNN(dim=dim, update_coors=update_coors, **egnn_kwargs)
            for _ in range(layers)
        ])

    def forward(self, feats, coors, mask=None):
        # feats: [B, L_max, D], coors: [B, L_max, 3], mask: [B, L_max] (bool)
        for layer in self.layers:
            feats, coors = layer(feats, coors, mask=mask)
        return feats, coors

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

    def forward(self, inputs, targets):
        bce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
        p_t = torch.exp(-bce_loss)

        alpha_weight = self.alpha * targets + (1 - self.alpha) * (1 - targets)
        loss = alpha_weight * (1 - p_t) ** self.gamma * bce_loss

        if self.reduction == 'mean':
            return torch.mean(loss)
        elif self.reduction == 'sum':
            return torch.sum(loss)
        else:
            return loss
        
# ===================================== 主模型（完全版） ===========================================
class PeptideHLABindingPredictor(nn.Module):
    def __init__(
        self,
        phys_dim=20,               # 物化编码的输出维度（你定义的 PhysicochemicalEncoder）
        pep_dim=256,              # 统一后的 peptide 通道
        hla_dim=256,              # 统一后的 HLA 通道
        bilinear_dim=256,
        pseudo_seq_pos=None,      # 口袋位点（假定 0-based 且落在 [0,179]）
        device="cuda:0",
        loss_fn='bce',
        alpha=0.5,
        gamma=2.0,
        dropout=0.2,
        pos_weights=None
    ):
        super().__init__()
        self.device = device
        self.pep_dim = pep_dim
        self.hla_dim = hla_dim
        self.bilinear_dim = bilinear_dim
        self.alpha = alpha
        self.gamma = gamma
        self.dropout = dropout
        if loss_fn == 'bce':
            self.loss_fn = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([pos_weights]) if pos_weights is not None else None)
        elif loss_fn == 'focal':
            self.loss_fn = FocalLoss(alpha=alpha, gamma=gamma)
        else:
            raise ValueError(f"Unknown loss function: {loss_fn}")

        self.se3_model = StackedEGNN(
            dim=17, layers=3
        )
        
        self.max_pep_len = 20  
        self.max_hla_len = 180 
        
        self.pep_pos_embed = nn.Parameter(torch.randn(self.max_pep_len, pep_dim))
        self.hla_pos_embed = nn.Parameter(torch.randn(self.max_hla_len, hla_dim))

        # —— 分支投影到统一维度（逐残基）——
        # peptide 分支（Physicochem -> pep_dim, ESM2(1280) -> pep_dim）
        self.proj_pep_phys = ResidueProjector(in_dim=phys_dim, out_dim=pep_dim)  # 你的 PhysEnc 输出维设成 pep_dim
        self.proj_pep_esm  = ResidueProjector(in_dim=1280, out_dim=pep_dim)

        # HLA 分支（Physicochem -> hla_dim, ESM2(1280) -> hla_dim, Struct(17/或se3_out) -> hla_dim）
        self.proj_hla_phys = ResidueProjector(in_dim=phys_dim, out_dim=hla_dim)  # 你的 PhysEnc 输出维设成 hla_dim
        self.proj_hla_esm  = ResidueProjector(in_dim=1280, out_dim=hla_dim)
        self.proj_hla_se3  = ResidueProjector(in_dim=17, out_dim=hla_dim)  # 让 se3_model 输出维就是 hla_dim

        # —— 门控融合（逐残基）——
        self.gate_pep = ResidueDoubleFusion(pep_dim)   # pep_phys × pep_esm
        self.gate_hla = ResidueTripleFusion(hla_dim)  # hla_phys × hla_esm × hla_struct
        
        d_model = self.pep_dim
        n_heads = 8
        
        # 1. 用于 "Peptide 查询 HLA" (pep_q_hla_kv)
        self.cross_attn_pep_hla = nn.MultiheadAttention(
            embed_dim=d_model, 
            num_heads=n_heads, 
            dropout=self.dropout, 
            batch_first=True
        )
        self.norm_cross_pep = nn.LayerNorm(d_model)

        # 2. 用于 "HLA 查询 Peptide" (hla_q_pep_kv)
        self.cross_attn_hla_pep = nn.MultiheadAttention(
            embed_dim=d_model, 
            num_heads=n_heads, 
            dropout=self.dropout, 
            batch_first=True
        )
        self.norm_cross_hla = nn.LayerNorm(d_model)

        # —— 交互模块（Bilinear attention map）——
        self.bi_attn = BANLayer(v_dim=pep_dim, q_dim=hla_dim, h_dim=bilinear_dim, h_out=4, k=3)

        # —— 头部 —— 
        self.head = nn.Sequential(
            nn.Linear(bilinear_dim, bilinear_dim),
            nn.ReLU(),
            nn.Linear(bilinear_dim, 1)
        )

        # —— 口袋位点 —— 
        if pseudo_seq_pos is None:
            pseudo_seq_pos = [i-2 for i in [7, 9, 24, 45, 59, 62, 63, 66, 67, 69, 70, 73, 74, 76, 77, 80, 81, 84, 95, 97, 99, 114, 116, 118, 143, 147, 150, 152, 156, 158, 159, 163, 167, 171]]
        self.register_buffer("contact_idx", torch.tensor(pseudo_seq_pos, dtype=torch.long))

        # --------------------------------------------
        # Transformer Encoders for peptide & HLA
        # --------------------------------------------
        encoder_layer_pep = TransformerEncoderLayer(
            d_model=pep_dim,        # 输入维度
            nhead=8,               # 注意力头数（可调）
            dim_feedforward=pep_dim*4,
            dropout=self.dropout,
            batch_first=True       # 输入形状 [B,L,D]
        )
        self.pep_encoder = TransformerEncoder(encoder_layer_pep, num_layers=2)  # 可以调整层数

        encoder_layer_hla = TransformerEncoderLayer(
            d_model=hla_dim,
            nhead=8,
            dim_feedforward=hla_dim*4,
            dropout=self.dropout,
            batch_first=True
        )
        self.hla_encoder = TransformerEncoder(encoder_layer_hla, num_layers=1)

    # -------------------------- 工具：把 list of [L,D] pad 成 [B,L_max,D] --------------------------
    def _pad_stack(self, tensors, L_max=None):
        Ls = [t.shape[0] for t in tensors]
        if L_max is None: L_max = max(Ls)
        D = tensors[0].shape[-1]
        B = len(tensors)
        out = tensors[0].new_zeros((B, L_max, D))
        mask = torch.zeros(B, L_max, dtype=torch.bool, device=out.device)
        for i, t in enumerate(tensors):
            L = t.shape[0]
            out[i, :L] = t
            mask[i, :L] = True
        return out, mask  # [B,L_max,D], [B,L_max]

    # ----------------------------------- 口袋掩码 --------------------------------------
    
    def _mask_to_pockets(self, hla_feat):
        """
        从 HLA 特征中只保留 pocket 位点，返回：
        - hla_pocket: [B, n_pocket, D]
        - pocket_mask: [B, n_pocket] (全部 True)
        """
        B, L, D = hla_feat.shape

        # ensure idx in [0, L-1]
        idx = self.contact_idx.clamp(min=0, max=L-1)
        # gather pocket features
        hla_pocket = hla_feat[:, idx, :]     # [B, n_pocket, D]

        return hla_pocket
    
    def add_positional_encoding(self, x, pos_embed):
        """
        x: [B, L, D]
        pos_embed: [L_max, D]
        """
        B, L, D = x.shape
        # 截取前 L 个位置编码
        pe = pos_embed[:L, :].unsqueeze(0).expand(B, -1, -1)  # [B, L, D]
        return x + pe

    def forward(self, batch):
        # take batch from DataLoader
        pep_phys = batch['pep_phys'].to(self.device, non_blocking=True)
        pep_esm  = batch['pep_esm'].to(self.device, non_blocking=True)
        hla_phys = batch['hla_phys'].to(self.device, non_blocking=True)
        hla_esm  = batch['hla_esm'].to(self.device, non_blocking=True)
        hla_struct = batch['hla_struct'].to(self.device, non_blocking=True)
        hla_coord  = batch['hla_coord'].to(self.device, non_blocking=True)
        labels = batch['label'].to(self.device)

        # 1) peptide 物化 + ESM2 → gate 融合
        pep_phys = self.proj_pep_phys(pep_phys)
        pep_esm  = self.proj_pep_esm(pep_esm)
        pep_feat = self.gate_pep(pep_phys, pep_esm)   # [B, Lp, D]
        
        pep_feat = self.add_positional_encoding(pep_feat, self.pep_pos_embed)
        pep_feat = self.pep_encoder(pep_feat, src_key_padding_mask=~batch['pep_mask'].to(self.device, non_blocking=True))
        
        # 2) HLA 物化 + ESM2 + 结构 → SE3 → gate 融合
        hla_phys = self.proj_hla_phys(hla_phys)
        hla_esm  = self.proj_hla_esm(hla_esm)
        # hla_struct 是 [B, 180, 17]，先过 SE3
        hla_se3 = self.se3_model(hla_struct, hla_coord, None)[0]  # [B, 180, 17]
        hla_se3 = self.proj_hla_se3(hla_se3) # →256
        hla_feat = self.gate_hla(hla_phys, hla_esm, hla_se3)
        
        hla_feat = self.add_positional_encoding(hla_feat, self.hla_pos_embed)
        hla_feat = self.hla_encoder(hla_feat)
        
        # cross attention for pep
        pep_feat_cross, _ = self.cross_attn_pep_hla(
            query=pep_feat,
            key=hla_feat,
            value=hla_feat,
            key_padding_mask=None
        )

        # cross attention for hla
        hla_feat_cross, _ = self.cross_attn_hla_pep(
            query=hla_feat,
            key=pep_feat,
            value=pep_feat,
            key_padding_mask=~batch['pep_mask'].to(self.device, non_blocking=True)
        )
        
        pep_feat_updated = self.norm_cross_pep(pep_feat + pep_feat_cross)
        hla_feat_updated = self.norm_cross_hla(hla_feat + hla_feat_cross)

        # 3) mask HLA 口袋位点
        hla_pocket = self._mask_to_pockets(hla_feat_updated)

        # 4) bilinear attention
        fused_vec, attn = self.bi_attn(
            pep_feat_updated,
            hla_pocket,
            v_mask=batch['pep_mask'].to(self.device, non_blocking=True),
            q_mask=None
        )
        logits = self.head(fused_vec).squeeze(-1)
        
        probs = torch.sigmoid(logits).detach().cpu().numpy()

        binding_loss = self.loss_fn(logits, labels.float())

        return probs, binding_loss, attn.detach().cpu().numpy().sum(axis=1), fused_vec.detach().cpu().numpy()

    # -------------------------- 编码器复用接口（给 TCR-HLA 模型用） --------------------------
    def _pad_peptide(self, x, max_len):
        """Pad peptide feature tensor [1, L, D] to [1, max_len, D]."""
        B, L, D = x.shape
        if L < max_len:
            pad = x.new_zeros(B, max_len - L, D)
            return torch.cat([x, pad], dim=1)
        else:
            return x[:, :max_len, :]
        
    @torch.no_grad()
    def encode_peptide_hla(self, pep_id, pep_phys, pep_esm, hla_phys, hla_esm, hla_struct, hla_coord, max_pep_len):
        Lp = len(pep_id)
                
        pep_phys = self.proj_pep_phys(pep_phys)
        pep_esm  = self.proj_pep_esm(pep_esm)
        
        pep_phys = self._pad_peptide(pep_phys, max_pep_len)
        pep_esm  = self._pad_peptide(pep_esm,  max_pep_len)
    
        device = pep_phys.device
        pep_mask = torch.zeros(1, max_pep_len, dtype=torch.bool, device=device)
        pep_mask[0, :Lp] = True

        pep_feat = self.gate_pep(pep_phys, pep_esm)
        pep_feat = self.add_positional_encoding(pep_feat, self.pep_pos_embed)
        pep_feat = self.pep_encoder(pep_feat, src_key_padding_mask=~pep_mask)

        # 2) hla encoding
        hla_phys = self.proj_hla_phys(hla_phys)
        hla_esm  = self.proj_hla_esm(hla_esm)
        hla_se3  = self.se3_model(hla_struct, hla_coord, None)[0]
        hla_se3  = self.proj_hla_se3(hla_se3)
        hla_feat = self.gate_hla(hla_phys, hla_esm, hla_se3)
        hla_feat = self.add_positional_encoding(hla_feat, self.hla_pos_embed)
        hla_feat = self.hla_encoder(hla_feat)

        # --- 3a. Peptide (Q) 查询 HLA (K, V) ---
        pep_feat_cross, _ = self.cross_attn_pep_hla(
            query=pep_feat,
            key=hla_feat,
            value=hla_feat,
            key_padding_mask=None
        )
        pep_feat_updated = self.norm_cross_pep(pep_feat + pep_feat_cross)
        
        # --- 3b. HLA (Q) 查询 Peptide (K, V) ---
        hla_feat_cross, _ = self.cross_attn_hla_pep(
            query=hla_feat,
            key=pep_feat,
            value=pep_feat,
            key_padding_mask=~pep_mask
        )
        hla_feat_updated = self.norm_cross_hla(hla_feat + hla_feat_cross)

        return pep_feat_updated, hla_feat_updated

class TCRPeptideHLABindingPredictor(nn.Module):
    def __init__(
            self, 
            tcr_dim=256, 
            pep_dim=256, 
            hla_dim=256, 
            bilinear_dim=256, 
            loss_fn='bce',
            alpha=0.5,
            gamma=2.0,
            dropout=0.1,
            device='cuda:0',
            pos_weights=None
        ):
        super().__init__()
        
        # TCR α / β position embeddings
        self.max_tcra_len = 500
        self.max_tcrb_len = 500
        self.max_pep_len  = 20
        self.max_hla_len  = 180
        self.alpha = alpha
        self.gamma = gamma
        self.dropout = dropout

        if loss_fn == 'bce':
            self.loss_fn = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([pos_weights]) if pos_weights is not None else None)
        elif loss_fn == 'focal':
            self.loss_fn = FocalLoss(alpha=alpha, gamma=gamma)
        else:
            raise ValueError(f"Unknown loss function: {loss_fn}")

        self.tcra_pos_embed = nn.Parameter(torch.randn(self.max_tcra_len, tcr_dim))
        self.tcrb_pos_embed = nn.Parameter(torch.randn(self.max_tcrb_len, tcr_dim))
        self.pep_pos_embed  = nn.Parameter(torch.randn(self.max_pep_len, pep_dim))
        self.hla_pos_embed  = nn.Parameter(torch.randn(self.max_hla_len, hla_dim))

        self.device = device
        self.tcr_dim = tcr_dim
        self.pep_dim = pep_dim
        self.hla_dim = hla_dim
        self.bilinear_dim = bilinear_dim
        
        d_model = tcr_dim
        n_heads = 8

        self.cross_attn_tcra_pep = nn.MultiheadAttention(d_model, n_heads, dropout=self.dropout, batch_first=True)
        self.cross_attn_tcra_hla = nn.MultiheadAttention(d_model, n_heads, dropout=self.dropout, batch_first=True)
        self.cross_attn_tcrb_pep = nn.MultiheadAttention(d_model, n_heads, dropout=self.dropout, batch_first=True)
        self.cross_attn_tcrb_hla = nn.MultiheadAttention(d_model, n_heads, dropout=self.dropout, batch_first=True)
        self.norm_tcra_pep = nn.LayerNorm(d_model)
        self.norm_tcra_hla = nn.LayerNorm(d_model)
        self.norm_tcrb_pep = nn.LayerNorm(d_model)
        self.norm_tcrb_hla = nn.LayerNorm(d_model)

        # =======================
        # TCRα / TCRβ encoders
        # =======================
        def make_tcr_encoder():
            proj_phys = ResidueProjector(20, tcr_dim)
            proj_esm = ResidueProjector(1280, tcr_dim)
            proj_struct = ResidueProjector(17, tcr_dim)
            se3 = StackedEGNN(dim=17, layers=1)
            gate = ResidueTripleFusion(tcr_dim)
            encoder_layer = TransformerEncoderLayer(
                d_model=tcr_dim, nhead=8, dim_feedforward=tcr_dim*4, dropout=self.dropout, batch_first=True
            )
            encoder = TransformerEncoder(encoder_layer, num_layers=2)
            return nn.ModuleDict(dict(
                proj_phys=proj_phys, proj_esm=proj_esm, proj_struct=proj_struct,
                se3=se3, gate=gate, encoder=encoder
            ))

        self.tcra_enc = make_tcr_encoder()
        self.tcrb_enc = make_tcr_encoder()

        # =======================
        # Peptide encoder (phys + esm + structure)
        # =======================
        self.proj_pep_phys = ResidueProjector(20, pep_dim)
        self.proj_pep_esm  = ResidueProjector(1280, pep_dim)
        self.proj_pep_struct = ResidueProjector(17, pep_dim)
        self.pep_se3 = StackedEGNN(dim=17, layers=1)
        self.pep_gate = ResidueTripleFusion(pep_dim)
        pep_encoder_layer = TransformerEncoderLayer(
            d_model=pep_dim, nhead=8, dim_feedforward=pep_dim*4, dropout=self.dropout, batch_first=True
        )
        self.pep_encoder = TransformerEncoder(pep_encoder_layer, num_layers=2)

        # =======================
        # HLA encoder
        # =======================
        self.proj_hla_phys = ResidueProjector(20, hla_dim)
        self.proj_hla_esm = ResidueProjector(1280, hla_dim)
        self.proj_hla_struct = ResidueProjector(17, hla_dim)
        self.hla_se3 = StackedEGNN(dim=17, layers=1)
        self.hla_gate = ResidueTripleFusion(hla_dim)
        hla_encoder_layer = TransformerEncoderLayer(
            d_model=hla_dim, nhead=8, dim_feedforward=hla_dim*4, dropout=self.dropout, batch_first=True
        )
        self.hla_encoder = TransformerEncoder(hla_encoder_layer, num_layers=1)

        self.pep_gate_2 = ResidueDoubleFusion(pep_dim)
        self.hla_gate_2 = ResidueDoubleFusion(hla_dim)

        # =======================
        # Bilinear interactions
        # =======================
        self.bi_tcra_pep = BANLayer(tcr_dim, pep_dim, bilinear_dim, h_out=4, k=3)
        self.bi_tcrb_pep = BANLayer(tcr_dim, pep_dim, bilinear_dim, h_out=4, k=3)
        self.bi_tcra_hla = BANLayer(tcr_dim, hla_dim, bilinear_dim, h_out=4, k=3)
        self.bi_tcrb_hla = BANLayer(tcr_dim, hla_dim, bilinear_dim, h_out=4, k=3)

        # =======================
        # Head
        # =======================
        total_fused_dim = bilinear_dim * 4
        self.head = nn.Sequential(
            nn.Linear(total_fused_dim, bilinear_dim),
            nn.ReLU(),
            nn.Linear(bilinear_dim, 1)
        )

    def encode_tcr(self, x_phys, x_esm, x_struct, x_coord, x_mask, enc, pos_embed):
        phys = enc['proj_phys'](x_phys)
        esm  = enc['proj_esm'](x_esm)
        se3  = enc['se3'](x_struct, x_coord, None)[0]
        se3  = enc['proj_struct'](se3)
        feat = enc['gate'](phys, esm, se3)
        feat = self.add_positional_encoding(feat, pos_embed)
        feat = enc['encoder'](feat, src_key_padding_mask=~x_mask)
        return feat
    
    def add_positional_encoding(self, x, pos_embed):
        """
        x: [B, L, D]
        pos_embed: [L_max, D]
        """
        B, L, D = x.shape
        pe = pos_embed[:L, :].unsqueeze(0).expand(B, -1, -1)
        return x + pe
    
    # def _extract_cdr3_segment(self, tcr_feat, cdr3_start, cdr3_end):
    #     B, L, D = tcr_feat.shape
    #     device = tcr_feat.device

    #     max_len = (cdr3_end - cdr3_start + 1).max().item()

    #     # [max_len], 0..max_len-1
    #     rel_idx = torch.arange(max_len, device=device).unsqueeze(0).expand(B, -1)  # [B, max_len]
    #     # absolute index = start + rel_idx
    #     abs_idx = cdr3_start.unsqueeze(1) + rel_idx
    #     # clamp end
    #     abs_idx = abs_idx.clamp(0, L-1)

    #     # mask positions beyond end
    #     mask = rel_idx <= (cdr3_end - cdr3_start).unsqueeze(1)

    #     # gather
    #     # expand abs_idx to [B, max_len, D] for gather
    #     gather_idx = abs_idx.unsqueeze(-1).expand(-1, -1, D)
    #     out = torch.gather(tcr_feat, 1, gather_idx)  # [B, max_len, D]
        
    #     return out, mask
    
    def _extract_cdr3_segment(self, tcr_feat, cdr3_start, cdr3_end):
        """
        Extracts CDR3 embeddings and corresponding mask.
        tcr_feat: [B, L, D]
        cdr3_start, cdr3_end: [B]
        Returns:
            out:  [B, max_len, D]
            mask: [B, max_len]  (True = valid)
        """
        B, L, D = tcr_feat.shape
        device = tcr_feat.device

        # 每个样本的 cdr3 长度
        lens = (cdr3_end - cdr3_start).clamp(min=0)
        max_len = lens.max().item()

        rel_idx = torch.arange(max_len, device=device).unsqueeze(0).expand(B, -1)  # [B, max_len]
        abs_idx = cdr3_start.unsqueeze(1) + rel_idx  # [B, max_len]

        # mask: True 表示有效
        mask = rel_idx < lens.unsqueeze(1)  # 注意这里 "<" 就够了

        # 将超出范围的索引设为 0（任意有效索引都行，因为会被mask掉）
        abs_idx = torch.where(mask, abs_idx, torch.zeros_like(abs_idx))

        # gather
        gather_idx = abs_idx.unsqueeze(-1).expand(-1, -1, D)
        out = torch.gather(tcr_feat, 1, gather_idx)

        # 对 mask 为 False 的位置强制置零，避免无效 token 参与计算
        out = out * mask.unsqueeze(-1)

        return out, mask

    def forward(self, batch):
        # TCRα / TCRβ
        tcra_feat = self.encode_tcr(
            batch['tcra_phys'].to(self.device, non_blocking=True),
            batch['tcra_esm'].to(self.device, non_blocking=True),
            batch['tcra_struct'].to(self.device, non_blocking=True),
            batch['tcra_coord'].to(self.device, non_blocking=True),
            batch['tcra_mask'].to(self.device, non_blocking=True),
            self.tcra_enc,
            self.tcra_pos_embed
        )
        tcrb_feat = self.encode_tcr(
            batch['tcrb_phys'].to(self.device, non_blocking=True),
            batch['tcrb_esm'].to(self.device, non_blocking=True),
            batch['tcrb_struct'].to(self.device, non_blocking=True),
            batch['tcrb_coord'].to(self.device, non_blocking=True),
            batch['tcrb_mask'].to(self.device, non_blocking=True),
            self.tcrb_enc,
            self.tcrb_pos_embed
        )
        # peptide
        pep_phys = self.proj_pep_phys(batch['pep_phys'].to(self.device, non_blocking=True))
        pep_esm  = self.proj_pep_esm(batch['pep_esm'].to(self.device, non_blocking=True))
        pep_se3 = self.pep_se3(batch['pep_struct'].to(self.device, non_blocking=True), batch['pep_coord'].to(self.device, non_blocking=True), None)[0]
        pep_se3 = self.proj_pep_struct(pep_se3)
        pep_feat = self.pep_gate(pep_phys, pep_esm, pep_se3)
        pep_feat = self.add_positional_encoding(pep_feat, self.pep_pos_embed)
        pep_feat = self.pep_encoder(
            pep_feat,
            src_key_padding_mask=~batch['pep_mask'].to(self.device)
        )
        # hla
        hla_phys = self.proj_hla_phys(batch['hla_phys'].to(self.device, non_blocking=True))
        hla_esm  = self.proj_hla_esm(batch['hla_esm'].to(self.device, non_blocking=True))
        hla_se3 = self.hla_se3(batch['hla_struct'].to(self.device, non_blocking=True), batch['hla_coord'].to(self.device, non_blocking=True), None)[0]
        hla_se3 = self.proj_hla_struct(hla_se3)
        hla_feat = self.hla_gate(hla_phys, hla_esm, hla_se3)
        hla_feat = self.add_positional_encoding(hla_feat, self.hla_pos_embed)
        hla_feat = self.hla_encoder(hla_feat)

        if ('pep_feat_pretrain' in batch) and ('hla_feat_pretrain' in batch):
            pep_pretrain = batch['pep_feat_pretrain'].to(self.device, non_blocking=True)
            hla_pretrain = batch['hla_feat_pretrain'].to(self.device, non_blocking=True)

            # ---- 鲁棒的长度对齐 (裁剪到最小长度) ----
            Lp = pep_feat.shape[1]
            Lp_pretrain = pep_pretrain.shape[1]
            if Lp != Lp_pretrain:
                Lp_min = min(Lp, Lp_pretrain)
                pep_feat = pep_feat[:, :Lp_min, :]
                pep_pretrain = pep_pretrain[:, :Lp_min, :]

            Lh = hla_feat.shape[1]
            Lh_pretrain = hla_pretrain.shape[1]
            if Lh != Lh_pretrain:
                Lh_min = min(Lh, Lh_pretrain)
                hla_feat = hla_feat[:, :Lh_min, :]
                hla_pretrain = hla_pretrain[:, :Lh_min, :]

            # ---- Peptide gating ----
            pep_feat = self.pep_gate_2(pep_feat, pep_pretrain)
            # ---- HLA gating ----
            hla_feat = self.hla_gate_2(hla_feat, hla_pretrain)

        # TCRα CDR3 segment
        tcra_cdr3, cdr3a_mask = self._extract_cdr3_segment(
            tcra_feat,
            batch['cdr3a_start'].to(self.device, non_blocking=True),
            batch['cdr3a_end'].to(self.device, non_blocking=True)
        )

        # TCRβ CDR3 segment
        tcrb_cdr3, cdr3b_mask = self._extract_cdr3_segment(
            tcrb_feat,
            batch['cdr3b_start'].to(self.device, non_blocking=True),
            batch['cdr3b_end'].to(self.device, non_blocking=True)
        )
        
        # TCRα CDR3 ← Peptide
        tcra_cdr3_cross, _ = self.cross_attn_tcra_pep(
            query=tcra_cdr3,                  # [B, La_cdr3, D]
            key=pep_feat, value=pep_feat,     # [B, Lp, D]
            key_padding_mask=~batch['pep_mask'].to(self.device)
        )
        tcra_cdr3 = self.norm_tcra_pep(tcra_cdr3 + tcra_cdr3_cross)
        # 重新掩蔽 padding 的 CDR3 位置，防止无效 token 漏光
        tcra_cdr3 = tcra_cdr3 * cdr3a_mask.unsqueeze(-1)

        # TCRβ CDR3 ← Peptide
        tcrb_cdr3_cross, _ = self.cross_attn_tcrb_pep(
            query=tcrb_cdr3,
            key=pep_feat, value=pep_feat,
            key_padding_mask=~batch['pep_mask'].to(self.device)
        )
        tcrb_cdr3 = self.norm_tcrb_pep(tcrb_cdr3 + tcrb_cdr3_cross)
        tcrb_cdr3 = tcrb_cdr3 * cdr3b_mask.unsqueeze(-1)

        # ------------------ Cross-Attn：TCR 全序列 ↔ HLA（整条 TCR） ------------------
        # TCRα full ← HLA
        tcra_hla_cross, _ = self.cross_attn_tcra_hla(
            query=tcra_feat,                  # [B, La, D]
            key=hla_feat, value=hla_feat,     # [B, Lh, D]
            key_padding_mask=None
        )
        tcra_feat = self.norm_tcra_hla(tcra_feat + tcra_hla_cross)
        tcra_feat = tcra_feat * batch['tcra_mask'].to(self.device).unsqueeze(-1)

        # TCRβ full ← HLA
        tcrb_hla_cross, _ = self.cross_attn_tcrb_hla(
            query=tcrb_feat,
            key=hla_feat, value=hla_feat,
            key_padding_mask=None
        )
        tcrb_feat = self.norm_tcrb_hla(tcrb_feat + tcrb_hla_cross)
        tcrb_feat = tcrb_feat * batch['tcrb_mask'].to(self.device).unsqueeze(-1)

        # bilinear fusion
        vec_tcra_pep, attn_tcra_pep = self.bi_tcra_pep(tcra_cdr3, pep_feat, v_mask=cdr3a_mask, q_mask=batch['pep_mask'].to(self.device))
        vec_tcrb_pep, attn_tcrb_pep = self.bi_tcrb_pep(tcrb_cdr3, pep_feat, v_mask=cdr3b_mask, q_mask=batch['pep_mask'].to(self.device))
        vec_tcra_hla, attn_tcra_hla = self.bi_tcra_hla(tcra_feat, hla_feat, v_mask=batch['tcra_mask'].to(self.device), q_mask=None)
        vec_tcrb_hla, attn_tcrb_hla = self.bi_tcrb_hla(tcrb_feat, hla_feat, v_mask=batch['tcrb_mask'].to(self.device), q_mask=None)
        
        attn_tcra_pep_small = attn_tcra_pep.sum(dim=1).float() 
        attn_tcrb_pep_small = attn_tcrb_pep.sum(dim=1).float()
        attn_tcra_hla_small = attn_tcra_hla.sum(dim=1).float()
        attn_tcrb_hla_small = attn_tcrb_hla.sum(dim=1).float()

        attn_dict = {
            'attn_tcra_pep': attn_tcra_pep_small.detach().cpu().numpy(),
            'attn_tcrb_pep': attn_tcrb_pep_small.detach().cpu().numpy(),
            'attn_tcra_hla': attn_tcra_hla_small.detach().cpu().numpy(),
            'attn_tcrb_hla': attn_tcrb_hla_small.detach().cpu().numpy()
        }

        fused = torch.cat([vec_tcra_pep, vec_tcrb_pep, vec_tcra_hla, vec_tcrb_hla], dim=-1)
        logits = self.head(fused).squeeze(-1)
        
        labels = batch['label'].to(self.device)
        loss_binding = self.loss_fn(logits, labels.float())

        probs = torch.sigmoid(logits)
                
        return probs, loss_binding, pep_feat.detach().cpu().numpy(), attn_dict