# setup_pymol.py - PyMOL 自定义设置脚本

from pymol import cmd

def setup_rendering(bg_color='white', transparency=0.5, ray_shadow=False):
    """设置 PyMOL 的背景、透明度和渲染优化"""
    cmd.bg_color(bg_color)
    
    if ray_shadow:
        cmd.set('ray_shadows', 'on')
        cmd.set('ray_shadow', 0.2)
    else:
        cmd.set('ray_shadows', 'off')
        cmd.set('ray_shadow', 0)
    
    cmd.set('transparency_mode', 2)
    cmd.set('transparency', transparency)
    cmd.set('surface_quality', 1)
    
    # 移除了 surface_color_by_atom，因为某些版本不支持
    # 使用 color 命令可以达到类似效果
    
    print(f"✓ 渲染设置完成: 背景={bg_color}, 透明度={transparency}, 阴影={'开' if ray_shadow else '关'}")


def clean_structure(remove_water=True, remove_metals=True, remove_ions=True):
    """
    清理结构中的水分子、金属离子和其他离子
    
    参数:
        remove_water (bool): 是否删除水分子，默认 True
        remove_metals (bool): 是否删除金属离子，默认 True
        remove_ions (bool): 是否删除其他常见离子，默认 True
    """
    removed_items = []
    
    # 1. 删除水分子
    if remove_water:
        water_count = cmd.count_atoms('resn HOH+WAT+H2O+DOD+D2O')
        if water_count > 0:
            cmd.remove('resn HOH+WAT+H2O+DOD+D2O')
            removed_items.append(f'水分子 ({water_count} 个原子)')
    
    # 2. 删除常见金属离子
    if remove_metals:
        metals = [
            'NA', 'K', 'MG', 'CA', 'ZN', 'FE', 'MN', 'CU', 'NI', 'CO',
            'CD', 'HG', 'PB', 'AG', 'AU', 'PT', 'CR', 'MO', 'W', 'V'
        ]
        metal_selection = '+'.join(metals)
        metal_count = cmd.count_atoms(f'resn {metal_selection}')
        if metal_count > 0:
            cmd.remove(f'resn {metal_selection}')
            removed_items.append(f'金属离子 ({metal_count} 个原子)')
    
    # 3. 删除其他常见离子（氯离子、硫酸根等）
    if remove_ions:
        ions = ['CL', 'BR', 'I', 'F', 'SO4', 'PO4', 'NO3', 'NH4']
        ion_selection = '+'.join(ions)
        ion_count = cmd.count_atoms(f'resn {ion_selection}')
        if ion_count > 0:
            cmd.remove(f'resn {ion_selection}')
            removed_items.append(f'其他离子 ({ion_count} 个原子)')
    
    if removed_items:
        print("✓ 结构清理完成:")
        for item in removed_items:
            print(f"  - 已删除 {item}")
    else:
        print("✓ 未发现需要清理的分子")


def create_colored_objects(protein_color='orange', nucleic_color='white', ligand_color='red'):
    """创建蛋白质、核酸和小分子对象，并应用指定颜色"""
    if len(cmd.get_object_list()) == 0:
        print("⚠ 警告: 未检测到已加载的结构，请先加载 PDB 文件")
        return
    
    obj_created = []
    
    # 1. 创建蛋白质对象
    if cmd.count_atoms('polymer.protein') > 0:
        cmd.create('Protein_obj', 'polymer.protein')
        cmd.color(protein_color, 'Protein_obj')
        cmd.show('surface', 'Protein_obj')
        obj_created.append(f'Protein_obj ({protein_color}, {cmd.count_atoms("Protein_obj")} 原子)')
    
    # 2. 创建核酸对象
    if cmd.count_atoms('polymer.nucleic') > 0:
        cmd.create('Nucleic_obj', 'polymer.nucleic')
        cmd.color(nucleic_color, 'Nucleic_obj')
        cmd.show('surface', 'Nucleic_obj')
        obj_created.append(f'Nucleic_obj ({nucleic_color}, {cmd.count_atoms("Nucleic_obj")} 原子)')
    
    # 3. 创建小分子对象（排除水分子和离子）
    ligand_selection = 'organic and not (resn HOH+WAT+H2O+DOD+D2O)'
    if cmd.count_atoms(ligand_selection) > 0:
        cmd.create('Ligand_obj', ligand_selection)
        cmd.color(ligand_color, 'Ligand_obj')
        cmd.show('surface', 'Ligand_obj')
        obj_created.append(f'Ligand_obj ({ligand_color}, {cmd.count_atoms("Ligand_obj")} 原子)')
    
    # 缩放到所有对象
    cmd.zoom('all')
    
    if obj_created:
        print("✓ 对象创建完成:")
        for obj in obj_created:
            print(f"  - {obj}")
    else:
        print("⚠ 未创建任何对象，请检查结构类型")


def quick_setup(bg_color='white', transparency=0.5, 
                protein_color='orange', nucleic_color='white', ligand_color='red',
                clean=True):
    """
    一键完成所有设置：清理结构 + 渲染优化 + 对象创建与着色
    
    参数:
        bg_color (str): 背景颜色
        transparency (float): 透明度 (0-1)
        protein_color (str): 蛋白质颜色
        nucleic_color (str): 核酸颜色
        ligand_color (str): 小分子颜色
        clean (bool): 是否清理水分子和离子，默认 True
    
    使用示例:
        quick_setup()                                    # 默认设置+清理
        quick_setup(clean=False)                         # 不清理
        quick_setup('black', 0.3, 'cyan', 'yellow', 'magenta')
    """
    print("\n" + "="*60)
    print("开始 PyMOL 快速设置...")
    print("="*60)
    
    # 1. 清理结构（如果启用）
    if clean:
        clean_structure()
    
    # 2. 设置渲染
    setup_rendering(bg_color, transparency)
    
    # 3. 创建并着色对象
    create_colored_objects(protein_color, nucleic_color, ligand_color)
    
    print("\n" + "="*60)
    print("✓ PyMOL 快速设置完成！")
    print("="*60 + "\n")


# 注册函数到 PyMOL 的全局命名空间
cmd.extend('setup_rendering', setup_rendering)
cmd.extend('clean_structure', clean_structure)
cmd.extend('create_colored_objects', create_colored_objects)
cmd.extend('quick_setup', quick_setup)

print("\n" + "="*60)
print("PyMOL 自定义函数已加载")
print("\n可用命令:")
print("  setup_rendering          - 设置背景和渲染")
print("  clean_structure          - 清理水分子和离子")
print("  create_colored_objects   - 创建并着色对象")
print("  quick_setup              - 一键完成所有设置")
print("\n使用示例:")
print("  quick_setup()            # 默认设置（包含清理）")
print("  quick_setup(clean=False) # 不清理水和离子")
print("  clean_structure()        # 单独清理")
print("="*60 + "\n")