"""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_hydrogen_bonds()
  show_salt_bridges()
"""

from pymol.wizard import Wizard

cmd.load('https://bioinformatics.lt/ppi3d/downloads/example1/pdb_files/1_d1x9na2_1-d1x9na1_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 766+767+768+769+770+772+773+833')
cmd.select('B_interface', 'chain B and resi 351+374+375+376+377+378+379+449+451')
cmd.select('interface', 'A_interface or B_interface')
cmd.select('A_interface_H_bonds', 'chain A and resi 768')
cmd.select('B_interface_H_bonds', 'chain B and resi 351')
cmd.select('interface_H_bonds', 'A_interface_H_bonds or B_interface_H_bonds')
cmd.select('A_interface_salt_bridges', 'chain A and resi 768')
cmd.select('B_interface_salt_bridges', 'chain B and resi 351')
cmd.select('interface_salt_bridges', 'A_interface_salt_bridges or B_interface_salt_bridges')


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

    def show_hydrogen_bonds(self):
        visualize_interface('H_bonds', labels=True, distances=True)

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

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

    def get_panel(self):
        return [['1', 'PPI3D Interface', ''], ['2', 'Interface', 'cmd.get_wizard().show_interface()'], ['2', 'H-bonds', 'cmd.get_wizard().show_hydrogen_bonds()'], ['2', 'Salt bridges', 'cmd.get_wizard().show_salt_bridges()'], ['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_hydrogen_bonds', lambda: interface.show_hydrogen_bonds())
cmd.extend('show_salt_bridges', lambda: interface.show_salt_bridges())
cmd.extend('clear_labels', lambda: interface.clear_labels())