#!/usr/bin/env python
"""
Automated driver for modeling mutations, setting up AMBER systems, and simuating Stark shifts.

TODO

"""

#===============================================================================
# IMPORTS
#===============================================================================

import os
import os.path
import sys
import math
import time
import shutil

# Import OpenMM, units package, and application layer.
import simtk.unit as units
import simtk.openmm as openmm
import simtk.openmm.app as app

#===============================================================================
# MUTATIONS TO MAKE
#===============================================================================

# List of mutations to make.
mutation_list = [
    # Bosutinib in solvent.
    'solvent',
    # Wild-type.
    'WT',
    # Point mutants.
    'V281C',
    'M314H', 'M314Y', 'M314L',
    'V323C', 'V323S', 'V323T',
    'L325Y',
    'T338M', 'T338N', 'T338Q', 'T338V', 'T338I', 'T338F', 'T338L', 
    'Y340F', 'Y340H', 'Y340K', 
    'M341C', 'M341H',
    'A403C', 'A403S', 'A403T', 'A403V',
    # Double mutants.
    'V323C A403T', # (mimics EGFR ATP-binding site)
    'V323S A403T', # (mimics Her2 ATP-binding site)
    'T338M A403T', # (mimics PKA ATP-binding site)
    'T338M A403C', #
    'T338M A403S', #
    # Triple mutants.
    'T338M A403T V323T', 
    'M314L L325Y T338M', # kinases that bind bosutinib tightly but do not look like Src have these three mutations
    ]

target_simulation_length = 1000.0 * units.picoseconds # target simulation length

"""
# Extend simulations.
# List of mutations to make.
mutation_list = [
    # Bosutinib in solvent.
    'solvent',
    # Wild-type.
    'WT',
    # Point mutants.
    'T338M', 'T338Q',
    # Double mutants.
    'T338M A403C',
    # Triple mutants.
    'M314L L325Y T338M', # kinases that bind bosutinib tightly but do not look like Src have these three mutations
    ]
"""

target_simulation_length = 3000.0 * units.picoseconds # target simulation length

"""
# Extend simulations.
# List of mutations to make.
mutation_list = [
    # Wild-type.
    'WT',
    ]

target_simulation_length = 10000.0 * units.picoseconds # target simulation length
"""

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

if __name__ == '__main__':

    # Reference structure information.
    reference_model_filename = '../setup/modeller/src_tbosutinib_c4+d2_refine_waters_tls_38.modeller.pdb'
    chain = 'B'
    base_directory = 'src_tbosutinib_c4+d2_refine_waters_tls_38'

    # Generating mutations from mutated source file requires we first add restorative mutation T403A, making sure T403A A403T cancels!
    #reference_model_filename = '../setup/modeller/a403t_bosut_refine_nl5_30.modeller.pdb'
    #chain = 'B'
    #base_directory = 'a403t_bosut_refine_nl5_30'
    
    if not os.path.exists(base_directory):
        os.makedirs(base_directory)

    run_modeller = True
    run_amber = True
    run_stark = True

    # Use MPI, if available.
    print "Attempting to initialize MPI..."
    try:
        from mpi4py import MPI # MPI wrapper
        rank = MPI.COMM_WORLD.rank
        size = MPI.COMM_WORLD.size
        if MPI.COMM_WORLD.rank == 0:
            print "MPI initialized: rank %d / size %d" % (rank, size)
    except Exception as e:
        print e
        print "MPI unavailable; using serial execution..."
        rank = 0
        size = 1

    # Process mutations.
    for mutation in mutation_list[rank::size]:
        import os.path
        output_directory = os.path.join(base_directory, mutation)
        if not os.path.exists(output_directory):
            os.makedirs(output_directory)
        
        if run_modeller:
            # Make mutation.
            print "============================================================================================="
            print "Building mutation '%s'..." % (mutation)
            print "============================================================================================="
            initial_modules = sys.modules.keys()
            import mutate # TODO: Make sure MODELLER's NetCDF gets unloaded after this block to prevent version incompatabilities.
            receptor_pdb_filename = os.path.join(output_directory, 'receptor.pdb')
            if mutation == 'solvent':
                # Create empty receptor filename.
                outfile = open(receptor_pdb_filename, 'w')
                outfile.close()
                pass
            elif mutation == 'WT':
                # Copy initial structure without mutations.
                shutil.copy(reference_model_filename, receptor_pdb_filename)
            else:
                # Make mutation.
                mutate.mutate(reference_model_filename, mutation, chain=chain, verbose=False, outfile=receptor_pdb_filename)
            # Error checking.
            if not os.path.exists(receptor_pdb_filename):
                raise Exception("mutate stage failed for mutation '%s'" % mutation)
            # Clean up modules.
            for module in sys.modules.keys():
                if module not in initial_modules:
                    del sys.modules[module]

        if run_amber:
            # Create AMBER system.
            print "Creating AMBER system..."
            import amber
            receptor_pdb_filename = os.path.join(output_directory, 'receptor.pdb')
            leap_output_filename = os.path.join(output_directory, 'leap.out')
            crystallographic_waters_pdb_filename = '../setup/modeller/closewater.pdb'
            ligand_frcmod_filename = '../setup/antechamber/oechem/bosutinib.gaff.frcmod'
            ligand_mol2_filename = '../setup/antechamber/oechem/bosutinib.gaff.mol2'
            prefix = os.path.join(output_directory, 'leap')
            amber.create_system(receptor_pdb_filename, ligand_frcmod_filename, ligand_mol2_filename, prefix, otherpdb=crystallographic_waters_pdb_filename, autocharge=True, leap_output_filename=leap_output_filename, verbose=False)
            if not os.path.exists(prefix + '.complex.prmtop'):
                raise Exception("AmberTools stage failed for mutation '%s'" % mutation)

        if run_stark:
            # Run Stark simulation.
            print "Running Stark simulation..."
            from stark import readPdbAtoms, findAtoms, simulate_stark_shift
            prmtop_filename = os.path.join(output_directory, 'leap.complex.prmtop')
            inpcrd_filename = os.path.join(output_directory, 'leap.complex.inpcrd')
            pdb_filename = os.path.join(output_directory, 'leap.complex.pdb')
            alpha = 0.87 * (units.centimeters**-1) / (units.mega*units.volts/units.centimeter) # linear Stark tuning rate of bosutinib
            atoms = readPdbAtoms(pdb_filename)
            atom_indices = findAtoms(atoms, resName='DB8', name='C3') + findAtoms(atoms, resName='DB8', name='N1') # Stark vibrational group
            print atom_indices
            try:
                simulate_stark_shift(prmtop_filename, inpcrd_filename, pdb_filename, alpha, atom_indices, output_directory, target_simulation_length=target_simulation_length)
            except Exception as e:
                print str(e)    
            if not os.path.exists(os.path.join(output_directory, 'stark.nc')):
                print "Stark simulation stage failed for mutation '%s'" % mutation