macro_split/tests/test_ring_numbering.py

"""
测试环编号功能 - 验证原子编号是否固定
"""
import sys
sys.path.insert(0, '/home/zly/project/macro_split')

from rdkit import Chem
from rdkit.Chem import Draw, AllChem
from rdkit.Chem.Draw import rdMolDraw2D
from src.ring_visualization import (
    get_macrolactone_numbering,
    get_ring_atoms_by_size
)


def test_ring_numbering_consistency(smiles: str, ring_size: int = 16, num_tests: int = 5):
    """
    测试环编号的一致性 - 多次运行确保编号固定
    """
    print("=" * 70)
    print("测试环编号一致性")
    print("=" * 70)
    print(f"\nSMILES: {smiles[:80]}...")
    print(f"环大小: {ring_size}")
    print(f"测试次数: {num_tests}")

    # 解析分子
    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        print("❌ 无法解析SMILES")
        return False

    print(f"✓ 分子解析成功，共 {mol.GetNumAtoms()} 个原子")

    # 检测环大小
    ring_atoms = get_ring_atoms_by_size(mol, ring_size)
    if ring_atoms is None:
        for size in range(12, 21):
            ring_atoms = get_ring_atoms_by_size(mol, size)
            if ring_atoms:
                ring_size = size
                print(f"⚠️  使用检测到的{size}元环")
                break

    if ring_atoms is None:
        print("❌ 未找到12-20元环")
        return False

    print(f"✓ 找到{ring_size}元环，包含 {len(ring_atoms)} 个原子")

    # 多次测试编号一致性
    all_numberings = []
    all_carbonyl_carbons = []
    all_ester_oxygens = []

    for i in range(num_tests):
        result = get_macrolactone_numbering(mol, ring_size)
        ring_atoms_result, ring_numbering, ordered_atoms, carbonyl_carbon, ester_oxygen, (is_valid, reason) = result

        if not is_valid:
            print(f"❌ 第{i+1}次测试失败: {reason}")
            return False

        all_numberings.append(ring_numbering.copy())
        all_carbonyl_carbons.append(carbonyl_carbon)
        all_ester_oxygens.append(ester_oxygen)

    # 验证一致性
    print("\n" + "-" * 50)
    print("编号一致性检查:")
    print("-" * 50)

    is_consistent = True

    if len(set(all_carbonyl_carbons)) == 1:
        print(f"✓ 羰基碳位置一致: 原子索引 {all_carbonyl_carbons[0]}")
    else:
        print(f"❌ 羰基碳位置不一致: {all_carbonyl_carbons}")
        is_consistent = False

    if len(set(all_ester_oxygens)) == 1:
        print(f"✓ 酯氧位置一致: 原子索引 {all_ester_oxygens[0]}")
    else:
        print(f"❌ 酯氧位置不一致: {all_ester_oxygens}")
        is_consistent = False

    first_numbering = all_numberings[0]
    for i, numbering in enumerate(all_numberings[1:], 2):
        if numbering != first_numbering:
            print(f"❌ 第{i}次编号与第1次不一致")
            is_consistent = False
            break

    if is_consistent:
        print(f"✓ 所有{num_tests}次测试的编号完全一致")

    # 显示详细编号信息
    print("\n" + "-" * 50)
    print("环原子编号详情:")
    print("-" * 50)

    numbering = all_numberings[0]
    carbonyl_carbon = all_carbonyl_carbons[0]
    ester_oxygen = all_ester_oxygens[0]

    sorted_items = sorted(numbering.items(), key=lambda x: x[1])

    print(f"{'位置':<6} {'原子索引':<10} {'元素':<6} {'说明'}")
    print("-" * 40)

    for atom_idx, position in sorted_items:
        atom = mol.GetAtomWithIdx(atom_idx)
        symbol = atom.GetSymbol()
        note = ""
        if atom_idx == carbonyl_carbon:
            note = "← 羰基碳 (C=O)"
        elif atom_idx == ester_oxygen:
            note = "← 酯键氧"
        print(f"{position:<6} {atom_idx:<10} {symbol:<6} {note}")

    return is_consistent


def save_visualization(smiles: str, output_path: str, ring_size: int = 16):
    """保存分子可视化图片"""
    print("\n" + "=" * 70)
    print("保存可视化图片")
    print("=" * 70)

    mol = Chem.MolFromSmiles(smiles)
    if mol is None:
        print("❌ 无法解析SMILES")
        return

    for size in range(12, 21):
        ring_atoms = get_ring_atoms_by_size(mol, size)
        if ring_atoms:
            ring_size = size
            break

    result = get_macrolactone_numbering(mol, ring_size)
    ring_atoms, ring_numbering, ordered_atoms, carbonyl_carbon, ester_oxygen, (is_valid, reason) = result

    if not is_valid:
        print(f"❌ 无法获取编号: {reason}")
        return

    mol_copy = Chem.Mol(mol)
    AllChem.Compute2DCoords(mol_copy)

    for atom_idx in ring_atoms:
        if atom_idx in ring_numbering:
            atom = mol_copy.GetAtomWithIdx(atom_idx)
            atom.SetProp("atomNote", str(ring_numbering[atom_idx]))

    atom_colors = {}
    for atom_idx in ring_atoms:
        atom = mol.GetAtomWithIdx(atom_idx)
        symbol = atom.GetSymbol()

        if atom_idx == carbonyl_carbon:
            atom_colors[atom_idx] = (1.0, 0.6, 0.0)
        elif atom_idx == ester_oxygen:
            atom_colors[atom_idx] = (1.0, 0.4, 0.4)
        elif symbol == 'C':
            atom_colors[atom_idx] = (0.7, 0.85, 1.0)
        elif symbol == 'O':
            atom_colors[atom_idx] = (1.0, 0.7, 0.7)
        elif symbol == 'N':
            atom_colors[atom_idx] = (0.8, 0.7, 1.0)
        else:
            atom_colors[atom_idx] = (0.8, 1.0, 0.8)

    drawer = rdMolDraw2D.MolDraw2DSVG(1000, 1000)
    drawer.SetFontSize(14)
    drawer.DrawMolecule(mol_copy, highlightAtoms=list(ring_atoms), highlightAtomColors=atom_colors)
    drawer.FinishDrawing()
    svg = drawer.GetDrawingText()

    svg_path = output_path.replace('.png', '.svg')
    with open(svg_path, 'w', encoding='utf-8') as f:
        f.write(svg)
    print(f"✓ SVG已保存到: {svg_path}")

    try:
        drawer_png = rdMolDraw2D.MolDraw2DCairo(1000, 1000)
        drawer_png.SetFontSize(14)
        drawer_png.DrawMolecule(mol_copy, highlightAtoms=list(ring_atoms), highlightAtomColors=atom_colors)
        drawer_png.FinishDrawing()
        drawer_png.WriteDrawingText(output_path)
        print(f"✓ PNG已保存到: {output_path}")
    except Exception as e:
        print(f"⚠️  PNG保存失败: {e}")

    print("\n颜色说明:")
    print("  橙色: 羰基碳 (位置1)")
    print("  红色: 酯键氧 (位置2)")
    print("  浅蓝色: 环上碳原子")


def main():
    smiles = "O[C@H]1[C@H]([C@H]([C@H](OC[C@@H]2[C@@H](CC)OC(C[C@H]([C@H](C)[C@H]([C@@H](CC=O)C[C@@H](C)C(/C=C/C(/C)=C/2)=O)O[C@H]2[C@@H]([C@H]([C@@H]([C@@H](C)O2)O[C@H]2C[C@](C)([C@@H]([C@@H](C)O2)O)O)[N@](C)C)O)O)=O)O[C@@H]1C)OC)OC"

    print("\n大环内酯环编号测试\n")
    is_consistent = test_ring_numbering_consistency(smiles, ring_size=16, num_tests=5)

    output_path = "/home/zly/project/macro_split/output/test_ring_numbering.png"
    save_visualization(smiles, output_path, ring_size=16)

    print("\n" + "=" * 70)
    print("测试总结")
    print("=" * 70)
    if is_consistent:
        print("✅ 所有测试通过！环原子编号是固定的。")
    else:
        print("❌ 测试失败：环原子编号不一致")

    return is_consistent


if __name__ == "__main__":
    success = main()
    sys.exit(0 if success else 1)