"""This is PyMOL script generated by PPI3D webserver.
Run it with PyMOL and explore the structure of protein-protein interface.

Usage:
  pymol script.py

Commands:
  show_interface()
  show_salt_bridges()
  show_disulphides()
  show_interface_ligands()
"""

from pymol.wizard import Wizard

cmd.load('https://bioinformatics.lt/ppi3d/downloads/example2/pdb_files/2_1s6v_D_1-1s6v_C_1.pdb')
chain_A_color = 'carbon' # PyMOL default
chain_B_color = 'cyan' # PyMOL default
chain_A_residues_color = 'lightorange'
chain_B_residues_color = 'lightmagenta'
cmd.set('seq_view', 'on')
cmd.set('label_position', '(2, 1, 2)')


def start():
    cmd.hide('all')
    cmd.show('cartoon')
    util.cbc()
    util.cnc()


def visualize(selection, how, color, labels):
    if selection in cmd.get_names('selections'):
        cmd.show(how, selection)
        cmd.color(color, selection)
        if labels:
            label_selection = '%s and name ca' % selection
            cmd.set('label_color', color, label_selection)
            cmd.label(label_selection, '"%s-%s" % (resn, resi)')
    else:
        return


def visualize_interface(what, how='sticks', labels=False, distances=False):
    global chain_A_residues_color
    global chain_B_residues_color
    start()
    suffix = ('_%s' % what) if what else ''
    a = 'A_interface%s' % suffix
    b = 'B_interface%s' % suffix
    visualize(a, how, chain_A_residues_color, labels)
    visualize(b, how, chain_B_residues_color, labels)
    if distances:
        cmd.delete('_distances')
        cmd.distance('_distances', a, b, mode=2)
    util.cnc()
    cmd.deselect()


cmd.select('A_interface', 'chain A and resi 12+13+16+17+28+70+72+73+79+81+82+83+86')
cmd.select('B_interface', 'chain B and resi 31+34+38+39+120+193+194+195+196+197+227+229+290')
cmd.select('interface', 'A_interface or B_interface')
cmd.select('A_interface_salt_bridges', 'chain A and resi 13')
cmd.select('B_interface_salt_bridges', 'chain B and resi 34')
cmd.select('interface_salt_bridges', 'A_interface_salt_bridges or B_interface_salt_bridges')
cmd.select('A_interface_disulphides', 'chain A and resi 81')
cmd.select('B_interface_disulphides', 'chain B and resi 197')
cmd.select('interface_disulphides', 'A_interface_disulphides or B_interface_disulphides')
cmd.select('A_interface_ligands', 'chain A and resi 104')
cmd.select('interface_ligands', 'A_interface_ligands')


class PPI3D_Interface(Wizard):
    def show_interface(self):
        visualize_interface('')

    def show_salt_bridges(self):
        visualize_interface('salt_bridges', labels=True, distances=True)

    def show_disulphides(self):
        visualize_interface('disulphides', labels=True)

    def show_interface_ligands(self):
        visualize_interface('ligands', 'spheres')

    def clear_labels(self):
        cmd.label('all', '')

    def get_panel(self):
        return [['1', 'PPI3D Interface', ''], ['2', 'Interface', 'cmd.get_wizard().show_interface()'], ['2', 'Salt bridges', 'cmd.get_wizard().show_salt_bridges()'], ['2', 'S-S bridges', 'cmd.get_wizard().show_disulphides()'], ['2', 'Interface ligands', 'cmd.get_wizard().show_interface_ligands()'], ['2', 'Clear labels', 'cmd.get_wizard().clear_labels()']]


start()
interface = PPI3D_Interface()
cmd.set_wizard(interface)
interface.show_interface()
cmd.extend('show_interface', lambda: interface.show_interface())
cmd.extend('show_salt_bridges', lambda: interface.show_salt_bridges())
cmd.extend('show_disulphides', lambda: interface.show_disulphides())
cmd.extend('show_interface_ligands', lambda: interface.show_interface_ligands())
cmd.extend('clear_labels', lambda: interface.clear_labels())