#! /usr/local/bin/python2.7
# =====================================================================
# pytleap version 1.3, April 2011
# copyright Novartis Institutes for Biomedical Research
# Basel, Switzerland
# Author: Romain M. Wolf
#======================================================================
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.

#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <http://www.gnu.org/licenses/>.
#=======================================================================
"""
python script calling tleap (and optionally antechamber) to generate
input files for AMBER modules (prmtop, crd, and pdb files);

organic ligands are treated via antechamber and 
must be supplied as MDL (SDF) file;

proteins and peptides are expected to be clean PDB files;

a complex between the ligand and the protein can be formed and the combined
AMBER prmtop, crd, and pdb files are also written out;

the main purpose is to generate the parameter-topology files for a subsequent
MMPB/SA-type computation, when separate prm files are required for the ligand,
the receptor, and the complex;

NEW February 2010: replaced MOPAC by sqm (works only for AMBER Tools 1.3 and later)
      NOTE: default settings for sqm are changed (see manual)
NEW April 2010: remove charge options Gasteiger and Mulliken (no good anyway)
    use only AM1-BCC by default (nothing else allowed)
NEW May 2010: protein-protein (or protein-peptide) complexes can also be
        generated now, the ligand must then be specified by the
        --pep option
NEW July 2010: adding --sspep option to make possible disulfides in pep also...
NEW November 2010: add some exception handling
"""

import os, sys
from optparse import OptionParser
from subprocess import Popen, PIPE
"""
Usage: pytleap.py [options] or -h (--help) for HELP

"""
#======================================================================
def tleap(cmd):
#======================================================================
  cmdline = "tleap -f %s > leap.output"%cmd
  os.system(cmdline)
#======================================================================
def antechamber(file, charge, qmet, fileformat):
#======================================================================
  """
  calls antechamber and parmchk to generate a MOL2 and an frcmod file
  for leap; cleans up some temporary files
  """
# for dmq, we use less severe convergence criteria than the default because the effect
# on resulting partial charges is negligible...
  sqmk = 'qm_theory=\'AM1\', grms_tol=0.05, tight_p_conv=1, scfconv=1.d-7, peptide_corr=1,'
  cmdline = 'antechamber -i %s '%(file)
  cmdline += '-fi %s -rn LIG -o %s.ac.mol2 -fo mol2 '%(fileformat, filename)
  cmdline += '-c bcc -nc %s -pf y -ek "%s"'%(charge, sqmk)
  cmdline += ' >/dev/null 2>&1'
  os.system(cmdline)
# make command line for parmchk
  cmdline = "parmchk -i %s.ac.mol2 -f mol2 -o %s.leap.frcmod"%(filename, filename)
# run parmchk to generate frcmod file for the organic ligand
  os.system(cmdline)

#*******************************************
# main calls
#*******************************************
if __name__ ==  "__main__":
  parser = OptionParser()
  parser.add_option("--prot", metavar = "FILE", dest = "prot", 
        help = "protein PDB file                       (no default)")
  parser.add_option("--pep", metavar = "FILE", dest = "pep", 
        help = "peptide PDB file                       (no default)")
  parser.add_option("--lig", metavar = "FILE", dest = "lig", 
        help = "ligand MDL (SDF) file                  (no default)")
  parser.add_option("--cplx", metavar = "FILE", dest = "cplx",
        help = "name for complex files                 (no default)")
  parser.add_option("--ppi", metavar = "FILE", dest = "ppi",
        help = "name for protein-peptide complex files (no default)")
  parser.add_option("--chrg", default = 0, metavar = "INTEGER", dest = "chrg",
        help = "formal charge on ligand                (default: 0)")
  parser.add_option("--rad", default = "mbondi", metavar = "STRING", dest = "rad",
        help = "radius type for GB                     (default: mbondi)")
  parser.add_option("--disul", metavar = "FILE", dest = "disul", 
        help = "file with S-S definitions in protein   (no default)")
  parser.add_option("--sspep", metavar = "FILE", dest = "sspep", 
        help = "file with S-S definitions in peptide   (no default)")
  parser.add_option("--pfrc", default = "ff03.r1", metavar = "STRING", dest = "pfrc", 
        help = "protein (peptide) force field          (default: ff03.r1)")
  parser.add_option("--lfrc", default = "gaff", metavar = "STRING", dest = "lfrc", 
        help = "ligand force field                     (default: gaff)")
  parser.add_option("--ctrl", default = "leap.cmd", metavar = "FILE", dest = "ctrl", 
        help = "leap command file name                 (default: leap.cmd)")

  if len(sys.argv) == 1:
    print("\n--------------------------------------------")
    print(" pytleap version 1.3 (April 2011)")
    print("--------------------------------------------")
    parser.print_help()
    sys.exit(-1)
    
  (opt, args) = parser.parse_args()

  if(opt.cplx) and (opt.ppi):
    print("You can't use --cplx and --ppi simultaneously! Bye!")
    sys.exit(-1)

# check if AMBERHOME is defined:
  amberhome = os.getenv("AMBERHOME")
  if not amberhome:
    print 'no AMBERHOME environment variable defined'
    print 'this means trouble...bye'

# the following executables must be in the path for pysolven to work:
  for prog in ['tleap', 'sqm', 'antechamber', 'parmchk']:
    try:
      Popen([prog, '-h'], stdin = PIPE, stdout = PIPE, stderr = PIPE)
    except OSError,e:
      print '...%s does not seem to be in the path...bye '%prog
      sys.exit(-1)

# check if specified input files exist:
  if opt.lig and not os.path.exists(opt.lig):
    print("...the ligand file %s cannot be found...bye"%opt.lig)
    sys.exit(-1)
  if opt.pep and not os.path.exists(opt.pep):
    print("...the peptide file %s cannot be found...bye"%opt.pep)
    sys.exit(-1)
  if opt.prot and not os.path.exists(opt.prot):
    print("...the protein file %s cannot be found...bye"%opt.prot)
    sys.exit(-1)

# write the leap control file
  leapcmd = opt.ctrl
  ctrl = open(leapcmd, 'w')
  frcprot = "leaprc.%s"%opt.pfrc
  frclig = "leaprc.%s"%opt.lfrc
  ctrl.write("source leaprc.%s\n"%opt.pfrc)
  ctrl.write("source leaprc.%s\n"%opt.lfrc)         
  ctrl.write("set default pbradii %s\n"%opt.rad)
# ligand preparation 
# NOTE: the mol2 file input is a hidden option because it does not necessarily
# work with just any mol2 file, so use very carefully and double-check results!!!
  if (opt.lig):
    filename = os.path.splitext(os.path.split(opt.lig)[1])[0]
    extension = os.path.splitext(os.path.split(opt.lig)[1])[1]
    if extension == '.mol2':
      fileformat='mol2'
    else:
      fileformat='mdl'

# call antechamber
    antechamber("%s"%opt.lig, opt.chrg, "bcc", fileformat)
# check if antechamber has succeeded by creating ...ac.mol2 file and
# if not, get out...
    if not os.path.exists('%s.ac.mol2'%filename):
      print(
        '''... antechamber seems to have failed...
... this could be due to a wrong formal charge for the ligand
... use the '--chrg' option correctly and try again...''')

    ctrl.write("lig = loadmol2 %s.ac.mol2\n"%filename)
    ctrl.write("frcmod = loadamberparams %s.leap.frcmod\n"%filename)
    ctrl.write("saveamberparm lig %s.leap.prm %s.leap.crd\n"%(filename, filename))
    ctrl.write("savepdb lig %s.leap.pdb\n"%filename)

# if ligand is a peptide, it is assigned the FF selected for proteins, i.e., no
# antechamber, no AM1-BCC charges, but plain protein; hence modified peptides
# will not work this way but must be used as 'ligand' with --lig...
  if (opt.pep):
    filename = os.path.splitext(os.path.split(opt.pep)[1])[0]
    ctrl.write("pep = loadpdb %s\n"%(opt.pep))

# we allow disulfide binds in the pep structure via the file read as opt.sspep
    if (opt.sspep):
      input = open(opt.sspep, 'r')
      for line in input.readlines():
        cys1 = line.split()[0]
        cys2 = line.split()[1]
        ctrl.write("bond pep.%s.SG pep.%s.SG\n"%(cys1, cys2))
    ctrl.write("saveamberparm pep %s.leap.prm %s.leap.crd\n"%(filename, filename))
    ctrl.write("savepdb pep %s.leap.pdb\n"%filename)

# protein (receptor) specification 
  if (opt.prot):
    filename = os.path.splitext(os.path.split(opt.prot)[1])[0]
    ctrl.write("prot = loadpdb %s\n"%(opt.prot))

# we treat disulfide bonds the old (secure) way, i.e., the user decides if and
# where they are... (and do not forget to rename the involved CYS to CYX!
    if (opt.disul):
      input = open(opt.disul, 'r')
      for line in input.readlines():
        cys1 = line.split()[0]
        cys2 = line.split()[1]
        ctrl.write("bond prot.%s.SG prot.%s.SG\n"%(cys1, cys2))
    ctrl.write("saveamberparm prot %s.leap.prm %s.leap.crd\n"%(filename, filename))
    ctrl.write("savepdb prot %s.leap.pdb\n"%filename)

# complex formation with organic ligand
  if (opt.cplx) and (opt.prot) and (opt.lig):
    ctrl.write("complex = combine {prot lig}\n")
    ctrl.write("saveamberparm complex %s.leap.prm %s.leap.crd\n" %(opt.cplx, opt.cplx))
    ctrl.write("savepdb complex %s.leap.pdb\n" %opt.cplx)     

# complex formation with peptide (or other protein)
  if (opt.ppi) and (opt.prot) and (opt.pep):
    ctrl.write("complex = combine {prot pep}\n")
    ctrl.write("saveamberparm complex %s.leap.prm %s.leap.crd\n" %(opt.ppi, opt.ppi))
    ctrl.write("savepdb complex %s.leap.pdb\n" %opt.ppi)      

  ctrl.write("quit\n")
  ctrl.close()
  tleap(leapcmd)