#!/usr/bin/python

#=============================================================================================
# Set up mol2 files for a halide scan.
#
# A sequence of mol2 files are written out where each proton is replaced with the specified
# halide atom.
#=============================================================================================

#=============================================================================================
# COPYRIGHT NOTICE
#
# Written by John D. Chodera <jchodera@berkeley.edu>, QB3-Berkeley Fellow.
#
# This file (only) is Copyright (c) 2009 The Regents of the University of California.
# All Rights Reserved.
#=============================================================================================

#=============================================================================================
# TODO
#=============================================================================================

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

# Python libraries
import os
import os.path
from optparse import OptionParser # for parsing of command line arguments
import sys
import commands # for running command-line utilities
import tempfile # for managing temporary directories

# OpenEye libraries
import openeye.oechem # use OpenEye OEChem toolkit for representing molecules and reading mol2 files

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

def write_file(filename, contents):
    """Write the specified contents to a file.

    ARGUMENTS
      filename (string) - the file to be written
      contents (string) - the contents of the file to be written

    """

    outfile = open(filename, 'w')

    if type(contents) == list:
        for line in contents:
            outfile.write(line)
    elif type(contents) == str:
        outfile.write(contents)
    else:
        raise "Type for 'contents' not supported: " + repr(type(contents))

    outfile.close()

    return

def read_file(filename):
    """Read contents of the specified file into a list of lines.

    ARGUMENTS
      filename (string) - the name of the file to be read

    RETURNS
      lines (list of strings) - the contents of the file, split by line
    """

    infile = open(filename, 'r')
    lines = infile.readlines()
    infile.close()

    return lines

def read_mol2_file(filename):
    """\
    Read a mol2 file as an OEChem molecule.
    Note that multi-molecule files are not supported.
    
    ARGUMENTS
      filename (string) - the name of the mol2 file to be read      
      
    RETURNS
      molecule (openeye.oechem.OEMol) - the molecule read in
      
    """

    # Open input mol2 file.
    input_molecule_stream = openeye.oechem.oemolistream()
    
    # Direct oemolistream to read all molecules in file into a multi-conformer molecule.
    input_molecule_stream.SetConfTest( openeye.oechem.OEIsomericConfTest() )    
    input_molecule_stream.open(filename)
    
    # Create a molecule.
    molecule = openeye.oechem.OEMol()
    
    # Read the molecule from file.
    openeye.oechem.OEReadMolecule(input_molecule_stream, molecule)
    
    # Return the moelecule.
    return molecule

def write_mol2_file(filename, molecule):
    """\
    Write an OEChem molecule to a mol2 file.
    
    ARGUMENTS
      filename (string) - the filename to write to.
      molecule (openeye.oechem.OEMol) - an OEChem molecule
    
    """
    
    # Create an output stream.
    output_molecule_stream = openeye.oechem.oemolostream()
    
    # Associate it with the given filename.
    output_molecule_stream.open(filename)
    
    # Write the molecule to the output stream.
    openeye.oechem.OEWriteMolecule(output_molecule_stream, molecule)
    
    # Close the output stream.
    output_molecule_stream.close()

    return

def formal_charge(molecule):
    """
    Report the net formal charge of a molecule.
    
    ARGUMENTS
       molecule (openeye.oechem.OEMol) - the molecule whose formal charge is to be determined

    RETURN VALUES
       formal_charge (integer) - the net formal charge of the molecule
       
    EXAMPLE
       net_charge = formal_charge(molecule)
   """

    # Create a copy of the molecule.
    molecule_copy = openeye.oechem.OEMol(molecule)

    # Assign formal charges.
    openeye.oechem.OEFormalPartialCharges(molecule_copy)

    # Compute net formal charge.
    formal_charge = int(round(openeye.oechem.OENetCharge(molecule_copy)))
    
    # return formal charge
    return formal_charge

def setup_complex(receptor_pdb_filename, ligand_mol2_filename, prefix, judgetypes=None, charge_model=None, formal_charge=0, protein_forcefield_leaprc='leaprc.ff96'):
    """
    Set up the specified receptor-ligand complex for AMBER simulation.

    ARGUMENTS

    receptor_pdb_filename (string) - PDB filename for input receptor
    ligand_mol2_filename (string) - Tripos format mol2 filename for input ligand
    prefix (string) - prefix for output (-complex.{prmtop,crd,pdb}, -ligand.{prmtop,crd,pdb})

    OPTIONAL ARGUMENTS

    pdb_filename (string) - if specified, will write a PDB file of the complex (default: None)
    judgetypes (int) - integer argument to Antechamber for judgetypes (default: None)
    charge_model (string) - (default: None)
    formal_charge (int) - (default: 0)
    protein_forcefield_leaprc (string) - default: leaprc.ff96

    """

    # DEBUG
    verbose = True
    show_warnings = True
    cleanup = False

    #
    # Parameterize small molecule with Antechamber.
    #

    # Create temporary working directory.
    import os, tempfile
    working_directory = tempfile.mkdtemp()
    old_directory = os.getcwd()
    os.chdir(working_directory)
    if verbose: print "Working directory is %(working_directory)s" % vars()

    # Run antechamber to assign GAFF atom types.
    import os.path, commands
    tripos_mol2_filename = os.path.join(old_directory, ligand_mol2_filename)
    gaff_mol2_filename = os.path.join(working_directory, 'gaff.mol2')   
    command = 'antechamber -i %(tripos_mol2_filename)s -fi mol2 -o %(gaff_mol2_filename)s -fo mol2' % vars()
    if judgetypes: command += ' -j %(judgetypes)d' % vars()
    if charge_model:
        command += ' -c %(charge_model)s -nc %(formal_charge)d' % vars()   
    if verbose: print command
    output = commands.getoutput(command)
    if verbose or (output.find('Warning')>=0): print output   
   
    # Generate frcmod file for additional GAFF parameters.
    frcmod_filename_tmp = os.path.join(working_directory, 'gaff.frcmod')
    commands.getoutput('parmchk -i %(gaff_mol2_filename)s -f mol2 -o %(frcmod_filename_tmp)s' % vars())

    #
    # Create AMBER topology/coordinate files for complex using LEaP.
    # 

    receptor_pdb_filename = os.path.join(old_directory, receptor_pdb_filename)

    # 
    leapscript = """\
# Load protein forcefield.
source %(protein_forcefield_leaprc)s

# Load ligand forcefield.
source leaprc.gaff

# Set GB radii to recommended values for OBC GBSA.
set default PBRadii mbondi2 
# Set GB radii to recommended values for GB/VI.
#set default PBRadii gbvi

# Load modifications to ligand forcefield for ligand.
mods = loadAmberParams %(frcmod_filename_tmp)s

# Load receptor.
receptor = loadpdb %(receptor_pdb_filename)s

# Load ligand.
ligand = loadmol2 %(gaff_mol2_filename)s

# Combine ligand and receptor
complex = combine { receptor ligand }

# DEBUG
check receptor
check ligand
check complex

# Save AMBER parmaeters.
saveAmberParm ligand ligand.prmtop ligand.crd
saveAmberParm complex complex.prmtop complex.crd

# QUIT
quit
""" % vars()
    leapin_filename = os.path.join(working_directory, 'leap.in')
    outfile = open(leapin_filename, 'w')
    outfile.write(leapscript)
    outfile.close()

    tleapout = commands.getoutput('tleap -f %(leapin_filename)s' % vars())
    if verbose: print tleapout
    tleapout = tleapout.split('\n')
    # Shop any warnings.
    if show_warnings:
        fnd = False
        for line in tleapout:
           tmp = line.upper()
           if tmp.find('WARNING')>-1: 
               print line
               fnd = True
           if tmp.find('ERROR')>-1: 
               print line
               fnd = True
        if fnd:
            print "Any LEaP warnings/errors are above."


    # Restore old directory.
    os.chdir(old_directory)   

    # Copy gromacs topology/coordinates to desired output files.
    commands.getoutput('cp %s %s' % (os.path.join(working_directory, 'complex.prmtop'), prefix + '-complex.prmtop'))
    commands.getoutput('cp %s %s' % (os.path.join(working_directory, 'complex.crd'), prefix + '-complex.crd'))
    commands.getoutput('cp %s %s' % (os.path.join(working_directory, 'ligand.prmtop'), prefix + '-ligand.prmtop'))
    commands.getoutput('cp %s %s' % (os.path.join(working_directory, 'ligand.crd'), prefix + '-ligand.crd'))

    # Create a PDB file.
    commands.getoutput('cat %s | ambpdb -p %s > %s' % (prefix + '-complex.crd', prefix + '-complex.prmtop', prefix + '-complex.pdb'))
    commands.getoutput('cat %s | ambpdb -p %s > %s' % (prefix + '-ligand.crd', prefix + '-ligand.prmtop', prefix + '-ligand.pdb'))
    
    # Clean up temporary files.
    os.chdir(old_directory)
    if cleanup:
        commands.getoutput('rm -r %s' % working_directory)
    else:
        print "Work done in %s..." % working_directory
    
    return

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

usage = """
  %prog will construct AMBER prmtop/crd files for a halide scan of a ligand-receptor binding free energy calculation

  USAGE

  %prog -r receptor.pdb { -n iupac_name | -i ligand.mol2 } -x halide_atom -o output_prefix

  EXAMPLES

  %prog -r example/receptor.pdb -n indole -x F -o example/complex
  %prog -r example/receptor.pdb -i example/indole.pdb -x Cl -o example/complex
  
"""

###############################################################################
# Process command-line options and check arguments.
#

# Create command-line argument option parser.
version = "%prog %__version__"

parser = OptionParser(usage=usage, version=version)


parser.add_option("-r", "--receptor", metavar='RECEPTOR_PDB_FILENAME',
                  action="store", type="string", dest='receptor_pdb_filename', default=None,                  
                  help="receptor PDB filename")
parser.add_option("-n", "--name", metavar='IUPAC_NAME',
                  action="store", type="string", dest='iupac_name', default=None,                  
                  help="an IUPAC name for the ligand molecule")
parser.add_option("-i", "--input", metavar='INPUT_MOL2_FILENAME',
                  action="store", type="string", dest='mol2_input_filename', default=None,                  
                  help="a Tripos mol2 file containing the ligand molecule to be halide scanned")
parser.add_option("-x", "--halide", metavar="HALIDE_ATOM_NAME",
                  action="store", dest="halide_atom_name", default='F', 
                  help="halide atom name to be used to replace protons (default: 'F')")
parser.add_option("-o", "--output", metavar="OUTPUT_PREFIX",
                  action="store", dest="output_prefix", default='F', 
                  help="output prefix")


# Parse command-line arguments.
(options,args) = parser.parse_args()

# Perform minimal error checking.
if not options.receptor_pdb_filename:
   parser.error("An input receptor PDB filename must be specified.")
if not (options.mol2_input_filename or options.iupac_name):
   parser.error("An input mol2/sdf/pdb file or IUPAC name must be specified.")
if not options.halide_atom_name:
   parser.error("A halide atom name must be specified.")   
if not options.output_prefix:
   parser.error("An output filename prefix must be specified.")

#if not os.path.exists(options.mol2_input_filename):
#   parser.error("Could not find specified mol2 input file %s" % options.mol2_input_filename)

if options.iupac_name:
    # Create molecule from IUPAC name
    import openeye.oeiupac
    molecule = openeye.oechem.OEMol()
    openeye.oeiupac.OEParseIUPACName(molecule, options.iupac_name)
    molecule.SetTitle(options.iupac_name)
elif options.mol2_input_filename:    
    # Read input mol2 file.
    molecule = read_mol2_file(options.mol2_input_filename)
    print "Molecule contains %d atoms." % molecule.NumAtoms()
    
    # TODO: Normalize.

# Determine what element to modify these atoms to.
atomic_number = openeye.oechem.OEGetAtomicNum(options.halide_atom_name)
print "Will change protons to element '%s', atomic number %d" % (options.halide_atom_name, atomic_number)

# Compute charges.
import openeye.oequacpac
noHydrogen = False
debug = True
openeye.oequacpac.OEAssignPartialCharges(molecule, openeye.oequacpac.OECharges_AM1BCC, noHydrogen, debug)

# Write file (normalizing it).
ligand_mol2_filename = options.output_prefix + '.mol2'
write_mol2_file(ligand_mol2_filename, molecule)

# Create AMBER prmtop/crd files for complex.
setup_complex(options.receptor_pdb_filename, ligand_mol2_filename, options.output_prefix)

# Make a list of atoms.
atoms = [ atom for atom in molecule.GetAtoms() ]
natoms = len(atoms)

# Make a list of the protons.
proton_indices = list()
for atom_index in range(natoms):
    if atoms[atom_index].GetAtomicNum() == 1:
        proton_indices.append(atom_index)
nprotons = len(proton_indices)
print "Molecule contains %d protons." % nprotons

# Modify one proton at a time.
for proton_index in proton_indices:
    # Copy molecule.
    modified_molecule = openeye.oechem.OEMol(molecule)

    # Get equivalent atom.
    atoms = [ atom for atom in modified_molecule.GetAtoms() ]
    atom = atoms[proton_index]

    # Change atom properties.
    atom.SetName(options.halide_atom_name)
    atom.SetAtomicNum(atomic_number)    

    # Construct modified prefix.
    modified_prefix = '%s-%d' % (options.output_prefix, proton_index)
        
    # Write file (normalizing it).
    ligand_mol2_filename = modified_prefix + '.mol2'
    write_mol2_file(ligand_mol2_filename, modified_molecule)

    # Assign charges.
    openeye.oequacpac.OEAssignPartialCharges(modified_molecule, openeye.oequacpac.OECharges_AM1BCC, noHydrogen, debug)

    # Write file with charges.
    write_mol2_file(ligand_mol2_filename, modified_molecule)

    # Create AMBER prmtop/crd files for complex.
    setup_complex(options.receptor_pdb_filename, ligand_mol2_filename, modified_prefix)