from pymol import cmd
import math

def calculate_gradient_transparency(obj, mole, num_segments=10, transparency_min=0.5, transparency_max=1.0):
    """
    设置基于距离的透明度梯度，并允许自定义透明度变化范围。
    
    参数:
    - obj: 对象名 (str)
    - mole: 小分子名 (str)
    - num_segments: 梯度分段数 (int)
    - transparency_min: 最小透明度 (float)
    - transparency_max: 最大透明度 (float)
    """
    # 检查 obj 是否存在
    if not cmd.count_atoms(obj):
        print(f"Error: The object '{obj}' does not exist or is empty. Please check your inputs.")
        return

    # 检查 mole 是否存在
    if not cmd.count_atoms(mole):
        print(f"Error: The object '{mole}' does not exist or is empty. Please check your inputs.")
        return

    # 获取复合物的边界框
    min_coords, max_coords = cmd.get_extent(obj)
    diagonal_length = math.sqrt(sum((max_coords[i] - min_coords[i])**2 for i in range(3)))

    # 手动计算 mole 的几何中心
    model = cmd.get_model(mole)
    x, y, z = 0, 0, 0
    for atom in model.atom:
        x += atom.coord[0]
        y += atom.coord[1]
        z += atom.coord[2]
    num_atoms = len(model.atom)
    center_coords = (x / num_atoms, y / num_atoms, z / num_atoms)

    # 打印计算信息
    print("\n[Gradient Transparency Calculation]")
    print(f"Bounding Box Min: {min_coords}")
    print(f"Bounding Box Max: {max_coords}")
    print(f"Diagonal Length: {diagonal_length:.3f} Å")
    print(f"Center of {mole}: ({center_coords[0]:.3f}, {center_coords[1]:.3f}, {center_coords[2]:.3f})\n")

    # 4. 计算梯度区间
    segment_length = diagonal_length / num_segments  # 每个梯度的长度
    print(f"Gradient Segments: {num_segments} (Length per Segment: {segment_length:.3f} Å)\n")

    # 打印透明度范围
    print(f"Transparency Range: {transparency_min} to {transparency_max}")

    # 创建透明度梯度
    for i in range(num_segments):
        start = i * segment_length
        end = (i + 1) * segment_length
        
        # 计算透明度值（线性插值）
        transparency = transparency_min + (transparency_max - transparency_min) * (i / (num_segments - 1))
        
        # 创建选择并设置透明度（略去重复代码）
        print(f"Segment {i+1}: {start:.3f}-{end:.3f} Å, Transparency: {transparency:.2f}")

    # 显示表面和卡通
    cmd.show("surface", obj)
    cmd.show("cartoon", obj)
    print("\n[Gradient Transparency Applied Successfully]\n")

# 调用函数（默认使用20个梯度）
calculate_gradient_transparency("all", "mole", num_segments=20)