import numpy as np
from rdkit import Chem
# from IsoSpecPy import IsoSpec
import IsoSpecPy as iso


def generate_isotopic_pattern(mol_formula, charge=1):
    """生成理论同位素模式"""
    # 初始化IsoSpec计算器
    # calculator = IsoSpec.IsoTotalProb(formula=mol_formula,
    #                                   charge=charge,
    #                                   absolute_accuracy=0.9)
    calculator = iso.IsoTotalProb(formula=mol_formula,
                                  charge=charge,
                                  absolute_accuracy=0.9)
    
    # 获取峰列表 (m/z, 强度)
    peaks = calculator.get_peaks()
    
    # 归一化强度
    intensities = np.array([p[1] for p in peaks])
    intensities /= np.max(intensities)
    
    return [(p[0], intensities[i]) for i, p in enumerate(peaks)]


def compare_isotopic_patterns(exp_pattern, theory_pattern, tolerance=0.01):
    """比较实验与理论同位素模式"""
    score = 0.0
    matched_peaks = 0
    
    for exp_mz, exp_int in exp_pattern:
        for theo_mz, theo_int in theory_pattern:
            if abs(exp_mz - theo_mz) <= tolerance:
                # 余弦相似度计算
                score += exp_int * theo_int
                matched_peaks += 1
                break
    
    # 归一化得分
    if matched_peaks > 0:
        score /= matched_peaks
    return score


# 示例用法
if __name__ == "__main__":
    # 实验数据 (示例值)
    experimental_data = [
        (151.063, 1.00),  # 基峰
        (152.066, 0.32),
        (153.069, 0.05)
    ]
    
    # 候选分子式
    candidate_formulas = ["C8H10N2O2", "C7H10N4O", "C9H10O3"]
    
    best_formula = None
    best_score = -1
    
    for formula in candidate_formulas:
        # 生成理论同位素模式
        theoretical_pattern = generate_isotopic_pattern(formula, charge=1)
        
        # 比较相似度
        score = compare_isotopic_patterns(experimental_data, theoretical_pattern)
        
        print(f"Formula: {formula} | Score: {score:.4f}")
        
        if score > best_score:
            best_score = score
            best_formula = formula
    
    print(f"\nBest match: {best_formula} (Score: {best_score:.4f})")