"""
Set up of solvated complexes using AmberTools LEaP.

@author John D. Chodera

"""

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

import sys
import os

#===============================================================================
# TEMPLATES
#===============================================================================

leap_templates = dict()

leap_templates['explicit'] = """
# Set up complex for explicit solvent simulation.

# Load AMBER ff99sb-ildn forcefield for protein.
source leaprc.%(protein_forcefield)s

# Load GAFF parameters for ligand.
source leaprc.gaff

# Load in protein (with all residues modeled in).
receptor = loadPdb "%(receptor_pdb_filename)s"

%(receptor_extra)s

# Add counterions
addions receptor Na+ %(ncations)d
addions receptor Cl- %(nanions)d

# Load ligand and parameters.
loadAmberParams "%(ligand_frcmod_filename)s"
ligand = loadMol2 "%(ligand_mol2_filename)s"

# Create complex.
complex = combine { receptor ligand }

# Check complex.
check complex

# Report on net charge.
charge ligand
charge receptor
charge complex

# Solvate ligand, receptor, and complex.
solvatebox ligand TIP3PBOX %(solvent_clearance)f %(closeness)f iso
solvatebox receptor TIP3PBOX %(solvent_clearance)f %(closeness)f iso
solvatebox complex TIP3PBOX %(solvent_clearance)f %(closeness)f iso

# Write parameters.
saveAmberParm ligand "%(prefix)s.ligand.prmtop" "%(prefix)s.ligand.inpcrd"
saveAmberParm receptor "%(prefix)s.receptor.prmtop" "%(prefix)s.receptor.inpcrd"
saveAmberParm complex "%(prefix)s.complex.prmtop" "%(prefix)s.complex.inpcrd"

# Exit
quit
"""

leap_templates['implicit'] = """
# Set up complex for GBSA simulation with OBC model.

# Set GB radii to recommended values for OBC.
set default PBRadii mbondi2

# Load AMBER ff99sb-ildn forcefield for protein.
source leaprc.%(protein_forcefield)s

# Load GAFF parameters for ligand.
source leaprc.gaff

# Load in protein (with all residues modeled in).
receptor = loadPdb %(receptor_filename)s

# Load ligand and parameters.
loadAmberParams %(ligand_frcmod_filename)s
ligand = loadMol2 %(ligand_mol2_filename)s

# Create complex.
complex = combine { receptor ligand }

# Check complex.
check complex

# Report on net charge.
charge ligand
charge receptor
charge complex

# Write parameters.
saveAmberParm ligand %(prefix)s.ligand.prmtop %(prefix)s.ligand.inpcrd
saveAmberParm receptor %(prefix)s.receptor.prmtop %(prefix)s.receptor.inpcrd
saveAmberParm complex %(prefix)s.complex.prmtop %(prefix)s.complex.inpcrd

# Exit
quit
"""

#===============================================================================
# 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.

    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 create_system(receptor_pdb_filename, ligand_frcmod_filename, ligand_mol2_filename, prefix, solvent='explicit', protein_forcefield='ff99SBildn', solvent_clearance=10.0, ncations=0, nanions=0, otherpdb=None):
    """
    Set up AMBER system using LEaP.

    ARGUMENTS

    receptor_pdb_filename (string) - source receptor filename (which must have crystallographic waters included already, if desired)
    ligand_frcmod_filename (string) - ligand frcmod filename
    ligand_mol2_filename (string) - ligand GAFF mol2 filename
    prefix (string) - prefix for writing out AMBER prmtop/inpcrd files (prefix-{receptor,ligand,complex}.{prmtop,inpcrd})

    OPTIONAL ARGUMENTS

    solvent (string) - either 'explicit' or 'implicit' (default: 'explicit')
    otherpdb (string) - if specified, add these (e.g. crystallographic waters) to receptor (default: None)

    TODO
    
    * Make this more general.
    * Automatically determine number of anions and cations based on system charge.
    
    """
    print "Constructing parameters..."

    parameters = dict()
    parameters['protein_forcefield'] = protein_forcefield
    parameters['solvent_clearance'] = solvent_clearance

    parameters['closeness'] = 0.75 # to permit crystallographic waters

    parameters['receptor_pdb_filename'] = receptor_pdb_filename
    parameters['ligand_frcmod_filename'] = ligand_frcmod_filename
    parameters['ligand_mol2_filename'] = ligand_mol2_filename

    parameters['prefix'] = prefix

    # TODO: Automate charge / counterion determination.
    parameters['ncations'] = ncations # number of cations to add
    parameters['nanions'] = nanions # number of anions to add

    # Add other PDB file to receptor if specified.
    parameters['receptor_extra'] = ''
    if otherpdb:
        parameters['receptor_extra'] = "other = loadPdb %s ; receptor = combine { receptor other }" % otherpdb
    print parameters
    
    # Write LEaP input to temporary file.
    print "Writing tleap.in..."
    tleap_input_filename = 'tleap.in'
    contents = leap_templates[solvent] % parameters
    write_file(tleap_input_filename, contents)
    
    # Run LEaP.
    print "Running LEaP..."
    import commands
    command = 'rm -f leap.log ; tleap -f %s' % tleap_input_filename
    output = commands.getoutput(command)
    print output

    # Generate PDB files.
    print "Generating PDB files..."
    for system in ['receptor', 'ligand', 'complex']:
        filename = '%s %s' % (prefix, system)
        command = 'cat "%(filename)s.inpcrd" | ambpdb -p "%(filename)s.prmtop" > "%(filename)s.pdb"' % vars()
        output = commands.getoutput(command)
        print output

    print "Done."

    return

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

if __name__ == '__main__':
    # Test on some data.

    receptor_pdb_filename = 'src-reference-model.pdb'
    crystallographic_waters_pdb_filename = '../modeller/closewater.pdb'
    ligand_frcmod_filename = '../antechamber/oechem/bosutinib.gaff.frcmod'
    ligand_mol2_filename = '../antechamber/oechem/bosutinib.gaff.mol2'
    prefix = 'test'
    ncations = 7
    nanions = 0

    create_system(receptor_pdb_filename, ligand_frcmod_filename, ligand_mol2_filename, prefix, ncations=ncations, nanions=nanions, otherpdb=crystallographic_waters_pdb_filename)