# Process structures with PyMOL.

# Disable GUI.
import __main__
__main__.pymol_argv = ['pymol','-qx'] # Pymol: quiet and no GUI
from time import sleep

# Launch PyMOL.
import pymol
pymol.finish_launching()

#===============================================================================
# SUBROUTINES
#===============================================================================

def create_image(filename='output.png', width=320, height=320):
    """
    Make a nice picture.

    ARGUMENTS

    filename (string) - output PNG filename (default: output.png)
    width (int) - width in pixels (default: 640)
    height (int) - height in pixels (default: 640)    

    """

    from pymol import cmd
    from pymol import util

    cmd.deselect()
    cmd.hide('all')
    cmd.show('cartoon', 'protein')
    cmd.show('sticks', 'ligand')
    cmd.show('spheres', 'waters')
    util.cbay('ligand')
    cmd.color('green', 'protein')
    cmd.set('surface_color', 'white')
    cmd.color('red', 'waters')
    cmd.show('surface', 'protein')
    cmd.set('transparency', 0.65)
    cmd.bg_color('white')
    cmd.set('antialias', 1)
    cmd.orient('system')
    cmd.zoom('system')
    cmd.viewport(width, height)
    cmd.ray(renderer='2')
    sleep(0.5) # short delay
    cmd.png(filename)

    return

#===============================================================================
# MAIN
#===============================================================================

from pymol import cmd

# Open a log.
log_filename = 'output.log'
cmd.log_open(log_filename)

# Load the X-ray structure.
cmd.load('xray-structures/src_tbosutinib_c4+d2_refine_waters_tls_38-seqres.pdb', 'pdb')

# Select protein, ligand, and nearby crystallographic waters to define system of interest.
cmd.select('protein', 'chain B')
cmd.select('ligand', 'chain C and resi 2')
cmd.select('waters', 'chain S within 6 of (protein or ligand)')
cmd.select('system', '(protein or ligand or waters)')

# Remove unused components.
cmd.remove('not system')

# deselect

# Export system components.
cmd.save('system.pdb', 'system')
cmd.save('protein.pdb', 'protein')
cmd.save('ligand.pdb', 'ligand')
cmd.save('waters.pdb', 'waters')

# Fetch template to use in modeling loops.
cmd.fetch('3lck', name='template')
cmd.remove('template and hetatm')

# Align template to protein.
cmd.align('template', 'protein', object='alignment')
cmd.save('template.pdb', 'template')
cmd.save('alignment.aln', 'alignment')

create_image()

# Close the log.
cmd.log_close()

# Quit.
cmd.quit()