dataset.py

import pandas as pd
import numpy as np
import itertools
import torch
from torch.utils.data import Dataset
import re
import json
from typing import Literal
import os
# import io
from rdkit import Chem
from rdkit.Chem import AllChem
from rdkit.Chem.Draw import rdMolDraw2D
# from PIL import Image
import torchvision.io as tvio
# import torchvision.transforms as tvt
import torchvision.transforms.v2.functional as tvtF

# --- 辅助函数 ---

# 定义20种常见氨基酸字母（按字母顺序）
AMINO_ACIDS = ['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 
               'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y']
AA_to_index = {aa: i for i, aa in enumerate(AMINO_ACIDS)}
valid_aa = set(AMINO_ACIDS)

def is_valid_sequence(seq):
    """
    判断序列是否只包含标准氨基酸字符（允许大写或小写，
    对于小写表示 D 型氨基酸也视为合法）
    """
    for ch in seq:
        if not ch.isalpha():
            return False
        if ch.upper() not in valid_aa:
            return False
    return True

def parse_mic(mic_str):
    """
    解析 MIC 数据，支持以下几种格式：
      1. 数字，例如 "5" -> 5.0 
      2. ">{数字}" 或 "≥{数字}"（例如 ">4" 或 "≥ 4"）→ 数值乘以 1.5
      3. 平均值±标准差，例如 "3.2 ± 0.4" → 取平均值 3.2
      4. 范围形式，例如 "2.0 - 4.0" → (2.0 + 4.0)/2

    注：符号与数字之间可能存在空格，大于等于符号为 "≥" 而非 ">="
    """
    if not isinstance(mic_str, str):
        return float(mic_str)
    
    mic_str = mic_str.strip()
    mic_str = re.sub(r'\s+', '', mic_str)
    
    # 匹配纯数字
    if re.fullmatch(r'\d+(\.\d+)?', mic_str):
        return float(mic_str)
    
    # 匹配 >{数字} 或 ≥{数字}
    m = re.fullmatch(r'[>≥](\d+(\.\d+)?)', mic_str)
    if m:
        num = float(m.group(1))
        return num * 1.5
    
    # 匹配 <{数字} 或 ≤{数字}
    m = re.fullmatch(r'[<≤](\d+(\.\d+)?)', mic_str)
    if m:
        num = float(m.group(1))
        return num * 0.75
    
    # 匹配 {数字}±{数字}
    m = re.fullmatch(r'(\d+(\.\d+)?)[±](\d+(\.\d+)?)', mic_str)
    if m:
        return float(m.group(1))
    
    # 匹配 {数字}-{数字}
    m = re.fullmatch(r'(\d+(\.\d+)?)-(\d+(\.\d+)?)', mic_str)
    if m:
        num1 = float(m.group(1))
        num2 = float(m.group(3))
        return (num1 + num2) / 2.0
    
    print(f"Warning: 无法解析 MIC 值 {mic_str}")
    return np.nan

def encode_sequence(seq, pad_length):
    """
    将多肽序列转换为固定大小 (pad_length, 21) 的张量：
      - 每个残基对应一行；
      - 第1列: 表示是否为 D 型氨基酸（若字符为小写，则置 1，否则为 0）；
      - 后20列: 20种常见氨基酸的独热编码（先转为大写匹配）。
    若序列长度小于 pad_length，则在末尾填充全 0 行。
    """
    n = len(seq)
    arr = np.zeros((pad_length, 21), dtype=np.float32)
    
    # 对实际序列部分进行编码
    for i, char in enumerate(seq):
        if i >= pad_length:
            break  # 超出部分不处理（数据集构造时已过滤掉长序列）
        if char.islower():
            d_indicator = 1.0
            aa = char.upper()
        else:
            d_indicator = 0.0
            aa = char
        arr[i, 0] = d_indicator
        if aa in AA_to_index:
            idx = AA_to_index[aa]
            arr[i, idx + 1] = 1.0
        else:
            print(f"Warning: 氨基酸 {aa} 不在标准列表中")
    return torch.tensor(arr)

def geometric_mean(values):
    """
    计算数值序列的几何平均值
    """
    log_vals = np.log(np.array(values))
    return float(np.exp(log_vals.mean()))

def process_label(ratio, task):
    """
    对比值 ratio 进行 log2 变换，并根据 task 参数返回最终标签：
     - task="reg": 返回 log₂比值，并转换为 np.float32；
     - task="cls": 根据 log₂比值进行分类：
                  如果 x <= -0.5 返回 1，
                  否则返回 0.
    若 ratio 非正，返回 np.nan。
    """
    if ratio <= 0:
        return np.nan
    ratio_log = np.log2(ratio)
    if task == "reg":
        return np.float32(ratio_log)
    elif task == "cls":
        if ratio_log < 0.:
            return 1
        else:
            return 0
    else:
        raise ValueError("未知的 task 类型，请使用 'reg' 或 'cls'")

# --- 数据预处理与构建数据集 ---

def load_data(xlsx_file, condition=None):
    """
    从 xlsx 文件中读取数据，将每个具体变种（同一原型-变种）对应的 MIC 值取几何平均，
    并按照原型分组。对于原型和变种序列，若存在非标准氨基酸或非字母字符，则过滤掉该行数据。

    返回：
      groups: dict，其中 key 为原型序列，
              value 为 dict，其 key 为变种序列（"SEQUENCE - D-type amino acid substitution"），
              value 为该变种所有 MIC 值的几何平均
    """
    df = pd.read_excel(xlsx_file)
    # df = df[df['TARGET ACTIVITY - ACTIVITY MEASURE VALUE'] != 'MBC']
    
    groups = {}
    for _, row in df.iterrows():
        orig = row["SEQUENCE - Original"]
        variant = row["SEQUENCE - D-type amino acid substitution"]
        mic_raw = row["TARGET ACTIVITY - CONCENTRATION"]
        
        # 过滤包含非标准氨基酸或非字母字符的序列（原型和变种均检查）
        if not (isinstance(orig, str) and is_valid_sequence(orig)):
            continue
        if not (isinstance(variant, str) and is_valid_sequence(variant)):
            continue
        
        mic_val = parse_mic(mic_raw)
        
        if orig not in groups:
            groups[orig] = {}
        if variant not in groups[orig]:
            groups[orig][variant] = []
        groups[orig][variant].append(mic_val)
    
    # 对每个变种计算几何平均（过滤掉 NaN 值）
    groups_avg = {}
    for orig, var_dict in groups.items():
        groups_avg[orig] = {}
        for variant, mic_list in var_dict.items():
            mic_list = [x for x in mic_list if not np.isnan(x)]
            if len(mic_list) == 0:
                continue
            groups_avg[orig][variant] = geometric_mean(mic_list)
    return groups_avg


def load_data_stability(xlsx_file, condition):
    """
    从 xlsx 文件中读取数据，将每个具体变种（同一原型-变种）对应的 MIC 值取几何平均，
    并按照原型分组。对于原型和变种序列，若存在非标准氨基酸或非字母字符，则过滤掉该行数据。

    返回：
      groups: dict，其中 key 为原型序列，
              value 为 dict，其 key 为变种序列（"SEQUENCE - D-type amino acid substitution"），
              value 为该变种所有 MIC 值的几何平均
    """
    map_dict = {
        '125fbs': '12.5% FBS',
        '25fbs': '25% FBS',
        'mhb': 'MHB',
        'nacl': '150mM NaCl'
    }
    df = pd.read_excel(xlsx_file)
    df = df[df['Condition'] == map_dict[condition]]
    
    groups = {}
    for _, row in df.iterrows():
        variant = row["SEQUENCE"]
        orig = variant.upper()
        mic_raw = row["Activity"]
        
        # 过滤包含非标准氨基酸或非字母字符的序列（原型和变种均检查）
        if not (isinstance(orig, str) and is_valid_sequence(orig)):
            continue
        if not (isinstance(variant, str) and is_valid_sequence(variant)):
            continue
        
        mic_val = parse_mic(mic_raw)
        
        if orig not in groups:
            groups[orig] = {}
        if variant not in groups[orig]:
            groups[orig][variant] = []
        groups[orig][variant].append(mic_val)
    
    # 对每个变种计算几何平均（过滤掉 NaN 值）
    groups_avg = {}
    for orig, var_dict in groups.items():
        groups_avg[orig] = {}
        for variant, mic_list in var_dict.items():
            mic_list = [x for x in mic_list if not np.isnan(x)]
            if len(mic_list) == 0:
                continue
            groups_avg[orig][variant] = geometric_mean(mic_list)
    return groups_avg

class PeptidePairDataset(Dataset):
    def __init__(self, mode=Literal['train', 'test', '125fbs', 'nacl', '25fbs', 'mhb'], pad_length=30, task="cls", 
                 include_reverse=False, include_self=False, one_way=False, gf=False) :
        """
        构建数据集：
          - 从 xlsx 文件中读取数据，并按照原型分组，
            同时过滤包含非标准氨基酸或非字母字符的行，以及变种序列长度超过 pad_length 的样本；
          - 对于同一原型下不同变种构成配对；
          - 参数 include_reverse: 是否启用正反组合（同时添加 (A, B) 和 (B, A)）；
          - 参数 include_self: 是否启用自组合（添加 (A, A)，标签为 log₂(1)=0）；
          - 参数 task: "reg" 表示回归任务（输出 32 位浮点数标签），"cls" 表示分类任务，
                      将 log₂比值转为整数标签，规则为：
                          log₂比值 ≤ -0.5 → 1，
                          log₂比值 ≥ 0.5 → 2，
                          -0.5 < log₂比值 < 0.5 → 0.
                          
        每个数据项返回：
          - 变种多肽序列编码后的张量，形状为 (pad_length, 21)
          - 另一个变种多肽序列编码后的张量，形状为 (pad_length, 21)
          - 标签：根据 task 不同分别为 32 位浮点数或整数
        """
        if mode == "train":
            loader = load_data
            xlsx_file = os.path.join(os.path.dirname(__file__), 'dataset', 'train.xlsx')
        elif mode in ["test", "r2_case", 'r2_case_', "125fbs", "nacl", "25fbs", "mhb"]:
            one_way = True
            if mode in ["test", "r2_case", 'r2_case_']:
                loader = load_data
                xlsx_file = os.path.join(os.path.dirname(__file__), 'dataset', f'{mode}.xlsx')
            else:
                loader = load_data_stability
                xlsx_file = os.path.join(os.path.dirname(__file__), 'dataset', 'stability.xlsx')
        else:
            raise ValueError("未知的 mode，请使用 'train' 或 'test'")

        self.data = []
        self.pad_length = pad_length
        self.task = task
        groups_avg = loader(xlsx_file, mode)
        if gf:
            gf_dict = torch.load(os.path.join(os.path.dirname(__file__), 'dataset', 'protbert.pth'))
        
        # 针对每个原型，过滤掉长度超过 pad_length 的变种
        for orig, variant_dict in groups_avg.items():
            # a = len(self.data)
            filtered_variants = {variant: mic for variant, mic in variant_dict.items() 
                                 if len(variant) <= pad_length}
            variants = list(filtered_variants.keys())
            n_variants = len(variants)
            if n_variants == 0:
                continue
            
            if gf:
                glob_feat = gf_dict[orig.upper()]

            # 若启用自组合，则添加 (A, A) 样本，标签为 process_label(1, task) → log2(1)=0（再分类也为 0）
            if include_self and (not one_way):
                for variant in variants:
                    encoded_seq = encode_sequence(variant, pad_length)
                    label = process_label(1.0, task)  # log2(1)=0
                    if gf:
                        self.data.append(((encoded_seq, encoded_seq, glob_feat), label))
                    else:
                        self.data.append(((encoded_seq, encoded_seq), label))
            
            # 添加不同变种之间的样本
            for i in [0] if one_way else range(n_variants):
                for j in range(i + 1, n_variants):
                    var1 = variants[i]
                    var2 = variants[j]
                    mic1 = filtered_variants[var1]
                    mic2 = filtered_variants[var2]
                    
                    # 正向组合: (var1, var2) 标签为 log₂(mic2/mic1)
                    ratio = mic2 / mic1 if mic1 != 0 else np.nan
                    label = process_label(ratio, task)
                    if np.isnan(label):
                        continue
                    encoded_var1 = encode_sequence(var1, pad_length)
                    encoded_var2 = encode_sequence(var2, pad_length)
                    if gf:
                        self.data.append(((encoded_var1, encoded_var2, glob_feat), label))
                    else:
                        self.data.append(((encoded_var1, encoded_var2), label))
                    
                    # 若启用正反组合，则添加 (var2, var1)
                    if include_reverse and (not one_way):
                        rev_ratio = mic1 / mic2 if mic2 != 0 else np.nan
                        rev_label = process_label(rev_ratio, task)
                        if gf:
                            self.data.append(((encoded_var2, encoded_var1, glob_feat), rev_label))
                        else:
                            self.data.append(((encoded_var2, encoded_var1), rev_label))
            # b = len(self.data)
            # print(f"{orig},{b - a}")
    
    def reg_sample_weight(self):
        y = []
        for _, label in self.data:
            y.append(label)
        y = np.array(y)
        mu = np.mean(y)
        sigma = np.std(y)
        p = 1 / (sigma * np.sqrt(2 * np.pi)) * np.exp(-((y - mu) ** 2) / (2 * sigma ** 2))
        
        # 如果未提供 C，则使用 p 的中位数作为基准常数
        C = np.median(p)
        epsilon = 1e-6
        
        # 使用对数转化计算采样权重： p 值越低权重越高
        weights = np.log(C / (p + epsilon))
        
        # 可选：对权重进行归一化处理，使得权重均值为1
        weights_normalized = weights / np.mean(weights)
        positive_weights = np.exp(weights_normalized)
        
        return torch.tensor(positive_weights, dtype=torch.float32)
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        return self.data[idx]


class PeptidePairPicDataset(Dataset):
    def __init__(self, mode=Literal['train', 'test', '125fbs', 'nacl', '25fbs', 'mhb'], pad_length=30, task="reg", 
                 include_reverse=False, include_self=False, one_way=False, gf=False,
                 side_enc=None, pcs=False, resize=None) :
        """
        构建数据集：
          - 从 xlsx 文件中读取数据，并按照原型分组，
            同时过滤包含非标准氨基酸或非字母字符的行，以及变种序列长度超过 pad_length 的样本；
          - 对于同一原型下不同变种构成配对；
          - 参数 include_reverse: 是否启用正反组合（同时添加 (A, B) 和 (B, A)）；
          - 参数 include_self: 是否启用自组合（添加 (A, A)，标签为 log₂(1)=0）；
          - 参数 task: "reg" 表示回归任务（输出 32 位浮点数标签），"cls" 表示分类任务，
                      将 log₂比值转为整数标签，规则为：
                          log₂比值 ≤ -0.5 → 1，
                          log₂比值 ≥ 0.5 → 2，
                          -0.5 < log₂比值 < 0.5 → 0.
                          
        每个数据项返回：
          - 变种多肽序列编码后的张量，形状为 (pad_length, 21)
          - 另一个变种多肽序列编码后的张量，形状为 (pad_length, 21)
          - 标签：根据 task 不同分别为 32 位浮点数或整数
        """
        if mode == "train":
            loader = load_data
            xlsx_file = os.path.join(os.path.dirname(__file__), 'dataset', 'train.xlsx')
        elif mode in ["test", "r2_case", 'r2_case_', "125fbs", "nacl", "25fbs", "mhb"]:
            one_way = True
            if mode in ["test", "r2_case", 'r2_case_']:
                loader = load_data
                xlsx_file = os.path.join(os.path.dirname(__file__), 'dataset', f'{mode}.xlsx')
            else:
                loader = load_data_stability
                xlsx_file = os.path.join(os.path.dirname(__file__), 'dataset', 'stability.xlsx')
        else:
            raise ValueError("未知的 mode，请使用 'train' 或 'test'")

        self.data = []
        self.pics = {}
        self.pad_length = pad_length
        self.task = task
        self.gf = gf
        self.side_enc = True if side_enc else False
        self.pcs = pcs
        self.resize = resize
        groups_avg = loader(xlsx_file, mode)
        if gf:
            gf_dict = torch.load(os.path.join(os.path.dirname(__file__), 'dataset', 'protbert.pth'))
        
        # 针对每个原型，过滤掉长度超过 pad_length 的变种
        for orig, variant_dict in groups_avg.items():
            # a = len(self.data)
            filtered_variants = {variant: mic for variant, mic in variant_dict.items() 
                                 if len(variant) <= pad_length}
            variants = list(filtered_variants.keys())
            for variant in variants:
                if self.pcs == 'mix' and variant == orig:
                    self.pics[variant] = self.read_img(variant, False)
                else:
                    self.pics[variant] = self.read_img(variant, self.pcs)
            n_variants = len(variants)
            if n_variants == 0:
                continue
            
            if gf:
                glob_feat = gf_dict[orig.upper()]

            # 若启用自组合，则添加 (A, A) 样本，标签为 process_label(1, task) → log2(1)=0（再分类也为 0）
            if include_self and (not one_way):
                for variant in variants:
                    label = process_label(1.0, task)  # log2(1)=0
                    if gf:
                        self.data.append((variant, variant, glob_feat, label))
                    else:
                        self.data.append((variant, variant, label))
            
            # 添加不同变种之间的样本
            for i in [0] if one_way else range(n_variants):
                for j in range(i + 1, n_variants):
                    var1 = variants[i]
                    var2 = variants[j]
                    mic1 = filtered_variants[var1]
                    mic2 = filtered_variants[var2]
                    
                    # 正向组合: (var1, var2) 标签为 log₂(mic2/mic1)
                    ratio = mic2 / mic1 if mic1 != 0 else np.nan
                    label = process_label(ratio, task)
                    if np.isnan(label):
                        continue
                    if gf:
                        self.data.append((var1, var2, glob_feat, label))
                    else:
                        self.data.append((var1, var2, label))
                    
                    # 若启用正反组合，则添加 (var2, var1)
                    if include_reverse and (not one_way):
                        rev_ratio = mic1 / mic2 if mic2 != 0 else np.nan
                        rev_label = process_label(rev_ratio, task)
                        if gf:
                            self.data.append((var2, var1, glob_feat, rev_label))
                        else:
                            self.data.append((var2, var1, rev_label))
            # b = len(self.data)
            # print(f"{orig},{b - a}")
    
    def reg_sample_weight(self):
        y = []
        for d in self.data:
            label = d[-1]
            y.append(label)
        y = np.array(y)
        mu = np.mean(y)
        sigma = np.std(y)
        p = 1 / (sigma * np.sqrt(2 * np.pi)) * np.exp(-((y - mu) ** 2) / (2 * sigma ** 2))
        
        # 如果未提供 C，则使用 p 的中位数作为基准常数
        C = np.median(p)
        epsilon = 1e-6
        
        # 使用对数转化计算采样权重： p 值越低权重越高
        weights = np.log(C / (p + epsilon))
        
        # 可选：对权重进行归一化处理，使得权重均值为1
        weights_normalized = weights / np.mean(weights)
        positive_weights = np.exp(weights_normalized)
        
        return torch.tensor(positive_weights, dtype=torch.float32)
    
    def read_img(self, peptide, pcs):
        image = draw_peptide(peptide, self.resize, pcs)
        return image
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        if self.gf:
            seq1, seq2, glob_feat, label = self.data[idx]
        else:
            seq1, seq2, label = self.data[idx]
        img1 = self.pics[seq1]
        img2 = self.pics[seq2]

        if self.side_enc:
            img1 = (img1, encode_sequence(seq1, self.pad_length))
            img2 = (img2, encode_sequence(seq2, self.pad_length))

        if self.gf:
            return (img1, img2, glob_feat), label
        else:
            return (img1, img2), label
        

class SimplePairClsDataset(Dataset):
    def __init__(self, pad_length=30, llm=False, ftr2=False, gf=False,
                 q_encoder=None, side_enc=None, pcs=False, resize=None):
        if llm:
            file_path = os.path.join(os.path.dirname(__file__), 'dataset', 'train_set_llm_aug.json')
        elif ftr2:
            file_path = os.path.join(os.path.dirname(__file__), 'dataset', 'finetune_for_r2_llm.json')
        else:
            file_path = os.path.join(os.path.dirname(__file__), 'dataset', 'train_set.json')
        with open(file_path, 'r', encoding='utf-8') as f:
            dataset = json.load(f)   

        self.data = []
        self.pics = {}
        self.pad_length = pad_length
        self.gf = gf
        self.q_encoder = q_encoder
        self.side_enc = True if side_enc else False
        self.pcs = pcs
        self.resize = resize
        if gf:
            self.gf_dict = torch.load(os.path.join(os.path.dirname(__file__), 'dataset', 'protbert.pth'))
        
        all_seqs = []
        for orig, variants in dataset.items():
            if len(orig) > pad_length:
                continue
            all_seqs.append(orig)
            for label in ["1", "0"]:
                for variant in variants[label]:
                    self.data.append((orig, variant, int(label)))
                    all_seqs.append(variant)
        if q_encoder in ['cnn', 'rn18']:
            for i in all_seqs:
                if self.pcs == 'mix' and i.isupper():
                    self.pics[i] = self.read_img(i, False)
                else:
                    self.pics[i] = self.read_img(i, self.pcs)
    
    def read_img(self, peptide, pcs):
        image = draw_peptide(peptide, self.resize, pcs)
        return image
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        seq1, seq2, label = self.data[idx]
        if self.q_encoder in ['cnn', 'rn18']:
            img1 = self.pics[seq1]
            img2 = self.pics[seq2]

            if self.side_enc:
                img1 = (img1, encode_sequence(seq1, self.pad_length))
                img2 = (img2, encode_sequence(seq2, self.pad_length))
        
        else:
            img1 = encode_sequence(seq1, self.pad_length)
            img2 = encode_sequence(seq2, self.pad_length)

        if self.gf:
            return (img1, img2, self.gf_dict[seq1]), label
        else:
            return (img1, img2), label


class PeptidePairCaseDataset(Dataset):
    def __init__(self, case:str ='r2', pad_length=30, gf=False):

        if case == 'r2':
            self.template = 'KWKIKWPVKWFKML'
        elif case == 'Indolicidin':
            self.template = 'ILPWKWPWWPWRR'
        elif case == 'Temporin-A':
            self.template = 'FLPLIGRVLSGIL'
        elif case == 'Melittin':
            self.template = 'GIGAVLKVLTTGLPALISWIKRKRQQ'
        elif case == 'Anoplin':
            self.template = 'GLLKRIKTLL'
        else:
            self.template = case.upper().strip()
        self.data = []
        self.pad_length = pad_length
        self.gf = gf

        if gf:
            self.glob_feat = torch.load(os.path.join(os.path.dirname(__file__), 'dataset', 'protbert.pth'))[self.template]
        
        pools = [(ch.upper(), ch.lower()) if ch != 'G' else (ch.upper(),) for ch in self.template]
        # 笛卡尔积，即所有组合
        self.variants = [''.join(chars) for chars in itertools.product(*pools)][1:]

        self.template_seq = encode_sequence(self.template, self.pad_length)
    
    def __len__(self):
        return len(self.variants)
    
    def __getitem__(self, idx):
        variant  = self.variants[idx]
        seq2, label = variant, variant
        enc_seq1 = self.template_seq
        enc_seq2 = encode_sequence(seq2, self.pad_length)

        if self.gf:
            return (enc_seq1, enc_seq2, self.glob_feat), label
        else:
            return (enc_seq1, enc_seq2), label



class PeptidePairPicCaseDataset(Dataset):
    def __init__(self, case:str ='r2', pad_length=30, side_enc=None, pcs=False, resize=None, gf=False):

        if case == 'r2':
            self.template = 'KWKIKWPVKWFKML'
        elif case == 'Indolicidin':
            self.template = 'ILPWKWPWWPWRR'
        elif case == 'Temporin-A':
            self.template = 'FLPLIGRVLSGIL'
        elif case == 'Melittin':
            self.template = 'GIGAVLKVLTTGLPALISWIKRKRQQ'
        elif case == 'Anoplin':
            self.template = 'GLLKRIKTLL'
        else:
            self.template = case.upper().strip()
        self.data = []
        self.pad_length = pad_length
        self.side_enc = True if side_enc else False
        self.pcs = pcs
        self.resize = resize
        self.gf = gf

        if gf:
            self.glob_feat = torch.load(os.path.join(os.path.dirname(__file__), 'dataset', 'protbert.pth'))[self.template]
        
        pools = [(ch.upper(), ch.lower()) if ch != 'G' else (ch.upper(),) for ch in self.template]
        # 笛卡尔积，即所有组合
        self.variants = [''.join(chars) for chars in itertools.product(*pools)][1:]

        self.template_pic = self.read_img(self.template)
        if self.side_enc:
            self.template_seq = encode_sequence(self.template, self.pad_length)
    
    def read_img(self, peptide):
        image = draw_peptide(peptide, self.resize, self.pcs)
        return image
    
    def __len__(self):
        return len(self.variants)
    
    def __getitem__(self, idx):
        variant  = self.variants[idx]
        seq2, label = variant, variant
        img1 = self.template_pic
        img2 = self.read_img(variant)

        if self.side_enc:
            img1 = (img1, self.template_seq)
            img2 = (img2, encode_sequence(seq2, self.pad_length))

        if self.gf:
            return (img1, img2, self.glob_feat), label
        else:
            return (img1, img2), label


aa_side = {
    "A": "C", "R": "CCCNC(N)=N", "N": "CC(=O)N", "D": "CC(=O)O", "C": "CS",
    "E": "CCC(=O)O", "Q": "CCC(=O)N", "G": "", "H": "Cc1cnc[nH]1", "I": "C(C)CC",
    "L": "CC(C)C", "K": "CCCCN", "M": "CCSC", "F": "Cc1ccccc1", "P": "C1CCN1",
    "S": "CO", "T": "C(C)O", "W": "Cc1c[nH]c2ccccc12", "Y": "Cc1ccc(O)cc1", "V": "C(C)C"
}

aa_tpl = {}
for aa, R in aa_side.items():
    for stereo, chir in (("L", "@"), ("D", "@@")):
        if aa == "G":  # Gly 没手性
            backbone = "N[C:{idx}]C"          # N-CA(带编号)-C
        else:
            backbone = f"N[C{chir}H:{'{idx}'}]({R})C"  # N-[C@H:idx](R)-C
        aa_tpl[f"{aa}_{stereo}"]       = backbone + "(=O)"     # 中间残基
        aa_tpl[f"{aa}_{stereo}_term"]  = backbone + "(=O)O"    # C 端

def build_peptide_smiles(seq: str) -> str:
    """
    给定单字母序列，返回 backbone 带 [atom_map] 的 SMILES。
    大写 = L 型， 小写 = D 型。编号 = 残基序号（1,2,3...） -> α-碳。
    """
    if not seq:
        return ""

    out = []
    n = len(seq)
    for i, aa in enumerate(seq, start=1):
        key = f"{aa.upper()}_{'L' if aa.isupper() else 'D'}"
        if i == n:
            key += "_term"
        out.append(aa_tpl[key].format(idx=i))
    return "".join(out)

protease_patterns = {
    'trypsin':       re.compile(r'(?<=[KR])(?!P)'),
    'chymotrypsin':  re.compile(r'(?<=[FYWL])(?!P)'),
    'elastase':      re.compile(r'(?<=[AVSGT])(?!P)'),
    'enterokinase':  re.compile(r'D{4}K(?=[^P])'),
    'caspase':       re.compile(r'(?<=D)(?=[GSA])'),
}

def draw_peptide(sequence, size=[768], pcs=False):
    """
    根据输入序列生成多肽结构图，并基于常见蛋白酶识别模式高亮酶切位点肽键（红色）。
    支持的酶及其正则模式（P1--P1'）：
      • trypsin:       (?<=[KR])(?!P)
      • chymotrypsin:  (?<=[FYWL])(?!P)
      • elastase:      (?<=[AVSGT])(?!P)
      • enterokinase:  D{4}K(?=[^P])
      • caspase:       (?<=D)(?=[GSA])
    """

    # # 1. 生成带 atom map 的 SMILES（现在序号标注在α-碳上）
    smiles = build_peptide_smiles(sequence)
    mol = Chem.MolFromSmiles(smiles)
    # if mol is None:
    #     raise ValueError("SMILES 解析失败，请检查输入序列和侧链字典。")
    AllChem.Compute2DCoords(mol)

    highlight_bonds = []
    bond_colors = {}
    
    # ----------------------------------------------------
    # 2. 先标 D 型残基：高亮与α-碳相连的键为蓝色
    d_positions = {i for i, aa in enumerate(sequence, start=1) if aa.islower()}

    for atom in mol.GetAtoms():
        if atom.GetAtomMapNum() in d_positions:
            # 这个atom就是α-碳，高亮与它相连的所有键
            for b in atom.GetBonds():
                idx = b.GetIdx()
                if idx not in highlight_bonds:
                    highlight_bonds.append(idx)
                bond_colors[idx] = (0.0, 0.0, 1.0)

    # ----------------------------------------------------
    # 3. 再标酶切键：红色（覆盖之前的蓝色）
    if pcs:
        cleavage_sites = set()
        for pat in protease_patterns.values():
            for m in pat.finditer(sequence):
                cut = m.end()  # 切在 cut 之后
                if 1 <= cut < len(sequence):
                    cleavage_sites.add(cut)

        for pos in cleavage_sites:
            # 先找 P1 残基的 α-C
            ca = next((a for a in mol.GetAtoms()
                       if a.GetAtomMapNum() == pos), None)
            if ca is None:
                continue

            # 找同残基的羧基碳 (sp², 含 O 双键)
            carbonyl_c = None
            for nb in ca.GetNeighbors():
                if nb.GetSymbol() != "C":
                    continue
                # 判断是否有 "=O"
                if any(bond.GetBondType() == Chem.BondType.DOUBLE and
                       o.GetSymbol() == "O"
                       for bond in nb.GetBonds()
                       for o in (bond.GetBeginAtom(), bond.GetEndAtom())):
                    carbonyl_c = nb
                    break
            if carbonyl_c is None:
                continue

            # 羧基碳连到的 N 就是下一残基的氮
            peptide_bond = None
            for b in carbonyl_c.GetBonds():
                o_atom = b.GetOtherAtom(carbonyl_c)
                if o_atom.GetSymbol() == "N":
                    peptide_bond = b
                    break
            if peptide_bond is None:
                continue

            bidx = peptide_bond.GetIdx()
            if bidx not in highlight_bonds:
                highlight_bonds.append(bidx)
            bond_colors[bidx] = (1.0, 0.0, 0.0)  # 红
    
    # 4. 设置画布大小
    if len(size) == 1:
        w = h = size[0]
    else:
        w, h = size

    # 5. MolDraw2DCairo 接收 highlightBondColors
    drawer = rdMolDraw2D.MolDraw2DCairo(w, h)
    # 你也可以通过 drawer.drawOptions() 调整一些样式：bond line width、atom font 等
    drawer.DrawMolecule(
        mol,
        highlightAtoms=[],
        highlightBonds=highlight_bonds,
        highlightAtomColors={},
        highlightBondColors=bond_colors
    )
    drawer.FinishDrawing()

    # 6. 把输出的 PNG bytes 转成 Tensor
    png_bytes = bytearray(drawer.GetDrawingText())
    byte_tensor = torch.frombuffer(png_bytes, dtype=torch.uint8)
    img = tvio.decode_png(byte_tensor, mode=tvio.ImageReadMode.RGB)       # [3, H, W], uint8
    img = tvtF.to_dtype(img, torch.float32)
    img = tvtF.normalize(img, [0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    return img

if __name__ == '__main__':
    # 假设 xlsx 文件路径为 "data.xlsx"
    # 设置 pad_length 为 50，同时启用正反组合和自组合
    pad_length = 30
    dataset = PeptidePairDataset('r2_case', pad_length, "cls", include_reverse=False, include_self=False, one_way=True)
    
    # 打印第一个数据项
    if len(dataset) > 0:
        (encoded_seq1, encoded_seq2), ratio = dataset[0]
        print("第一个样本：")
        print("变种1的编码张量形状：", encoded_seq1.shape)
        print("变种2的编码张量形状：", encoded_seq2.shape)
        print("标签比值（变种2/变种1）：", ratio)
        print(f"数据集大小：{len(dataset)}")
        label_pos = 0
        for (_, _), i in dataset:
            label_pos += i
        print(label_pos)

    else:
        print("未读入组合数据！")

    # # 测试 PeptidesDataset
    # pad_length = 30
    # dataset = PeptidesDataset(xlsx_file="./dataset/train.xlsx", pad_length=pad_length)
    # print(f"PeptidesDataset 样本总数: {len(dataset)}")
    # if len(dataset) > 0:
    #     encoded_seq, label = dataset[0]
    #     print("第一个样本：")
    #     print("多肽编码张量形状：", encoded_seq.shape)
    #     print("标签浓度值（几何平均后）：", label)
    # else:
    #     print("未读取到有效数据！")