"""
Reads in a 'PDBs' directory generated by SavePDBs.py and generates some
images that represent the conformations of molecules in each state.

THIS SCRIPT IS RUN FROM *WITHIN* PYMOL. While an external interface is
possible and desirable, it requires a specific install of PyMol on some
systems. This should be much more general.

Note that you probably want to play around with the parameters in
the RenderImage() function to get things looking right.

Example workflow to get some pretty images:
(1) Save PBSs from states that you care about using SavePDBs.py
(choose many confs, they become nice backgrounds)
(2) Open PyMol
(3) Run, in the command line:
> run GenerateStateFigure.py
> gen_state_fig <PDBs dir>, <aligned pdb>, <an output dir>

where "aligned pdb" is a PDB file of your molecule that is appropriately
aligned for visual effect. All structures will get aligned to this structure
by an RMSD minimization calculation, but this aligned/reference structure
will not appear in the final image.
"""

#import Trajectory
import sys
import os
import re
import glob
from pymol import cmd


def RenderImage(files, index, reference, OutDir):
    """This function issues commands to pymol. Alter these commands to change
    what the final output looks like. Note that 'reference' is a reference structure
    that allows you to align everything to that structure"""

    cmd.do('reinitialize')
    cmd.do('set_color lblue= [0.86 , 1.00 , 1.00]')

    for i,f in enumerate(files):
        cmd.load(f, 'Structure%d'%i)
    cmd.load(f, 'Overlay')
    cmd.load(reference, 'Reference')

    cmd.hide('all')
    cmd.do('dss')
    cmd.do('bg white')
    
    cmd.do('show cartoon, Overlay')
    cmd.fit('Overlay', 'Reference')

    cmd.fit('Structure0', 'Reference')
    for i in range(len(files)):
        cmd.do('show cartoon, Structure%d'%i)
        cmd.fit('Structure%d'%i, 'Structure0')
        cmd.do('color lblue, Structure%d'%i)
        cmd.do('set cartoon_transparency, 0.5, Structure%d'%i)
    cmd.do('spectrum count, rainbow, Overlay')
    cmd.do('translate [0,0,10], Overlay')

    cmd.do('ray')
    cmd.do('png %s/State%d'% (OutDir, index))

    return


def run(InDir, RefStruct, OutDir):

    # Check output paths
    if os.path.exists(OutDir):
        print "Error: Directory %s already exists! Exiting." % OutDir
        sys.exit(1)
    os.mkdir(OutDir)
    print "Created %s and wrote png files there" % OutDir

    # Parse directory
    ref_path = os.path.abspath(RefStruct)
    print "Looking in %s" % InDir

    states=[]
    files = glob.glob(InDir+"/State*-*.pdb")
    for f in files:
        state = int( re.findall("State(\d+)-\d+.pdb", f)[0] )
        if state not in states: states.append(state)

    for state in states:
        pdbs=glob.glob(InDir+'/State%d-*.pdb'%state)
        RenderImage(pdbs, state, ref_path, OutDir)

    print "Finished rendering all images."

    return


print """\nReads in a 'PDBs' directory generated by SavePDBs.py and generates a png
figure representing each state in each cluster. Also takes a reference structure,
which should be manually aligned and saved in some appealing conformation. All saved
figures will be aligned to this reference strucutre.

NOTE: This script can be VERY RESOURCE INTENSIVE since it has to ray-trace many images

usage syntax: 'gen_state_fig <pdb_input_dir>, <reference_structure>, <output_dir>'"""

cmd.extend('gen_state_fig', run)