from pymol import cmd
from pymol.cmd import util
import sys
from pymol_color_plugin import colorp


def HBond(ChA,ChB):
    name = f'Hbond'
    dist = cmd.distance(name, ChA,ChB, '3.5', '2')
    return name

def Pi_Pi_interaction(ChA,ChB):
    name = f'PI_PI'
    cmd.distance(name,ChA, ChB,'5.5','6')
    return name

def Cation_Pi_interaction(ChA,ChB):
    name = f'PI_Cation'
    cmd.distance(name, ChA, ChB, '5.5', '7')
    return name

def Salt_bridge_positive(ChA,ChB):
    name = f'SBP'

    selection_1 = f'{ChB} and ((resn LYS and name NZ) or (resn ARG and name NE+NH*))'
    selection_2 = f'{ChA} and resn ASP+GLU and name OD*+OE*'
    cmd.distance(name,selection_1,selection_2,'5.0','0')
    return name

def Salt_bridge_negative(ChA,ChB):
    name = f'SBN'
    selection_1 = f'{ChA} and ((resn LYS and name NZ) or (resn ARG and name NE+NH*))'
    selection_2 = f'{ChB} and resn ASP+GLU and name OD*+OE*'
    cmd.distance(name,selection_1,selection_2,'5.0','0')
    return name

# def T_stacking(ChA, ChB):
#     name = 'T_stacking'
#     # 添加适当的距离和选择标准
#     cmd.distance(name, ChA, ChB, '5.5', '2')
#     return name

# def Hydrophobic_interaction(ChA, ChB):
#     name = 'Hydrophobic'
#     # 添加适当的距离和选择标准
#     cmd.distance(name, ChA, ChB, '7.0', '2')
#     return name

# def VDW_interaction(ChA, ChB):
#     name = 'VDW'
#     # 添加适当的距离和选择标准
#     cmd.distance(name, ChA, ChB, '4.0', '2')
#     return name

def show_label_name(obj:str):
    cmd.label(f'name CA and {obj}', 'oneletter+resi')
    cmd.set('label_bg_color', -7, '', 0)
    cmd.set('label_size', 18, '', 0)
    cmd.set('label_font_id', '7')

def interact(obj:str='epitope', distance:int=5):
    cmd.show('stick',f'{obj}')
    cmd.set('stick_radius', '0.3', f'{obj}')
    colorp.color_mole(obj)
    #show_label_name(obj)
    cmd.select('nearby',f'byres {obj} around {distance}')
    cmd.show('stick','nearby')
    cmd.set('stick_radius', '0.1', 'nearby')
    cmd.set('cartoon_transparency', 0.7, selection='nearby')
    Hbond = HBond(obj, 'nearby')
    PI = Pi_Pi_interaction(obj, 'nearby')
    C_PI = Cation_Pi_interaction(obj, 'nearby')
    SBP = Salt_bridge_positive(obj, 'nearby')
    SBN = Salt_bridge_negative(obj, 'nearby')
    cmd.set('dash_color', 'yellow', Hbond)
    cmd.set('dash_color', 'blue', PI)
    cmd.set('dash_color', 'blue', C_PI)
    cmd.set('dash_color', 'orange', SBP)
    cmd.set('dash_color', 'orange', SBN)
    cmd.set('dash_radius', '0.07')
    cmd.set('dash_gap', '0.3')
    #cmd.remove('e. h and  neighbor(name C*)')
    cmd.zoom('Hbond')

def main():
    cmd.reinitialize()
    cmd.fetch('5sws',type='pdb')
    colorp.pretty()
    cmd.select('epitope','chain C and resi 1-9')
    cmd.zoom('nearby')
    interact()
    cmd.save(f'5sws_interaction.pse')