#! PYTHONEXE

"""
This program is will take a topology file and adjust the PB/GB
Radius set to match what the user puts on the command-line
"""

import sys
from optparse import OptionParser
from chemistry.amber.readparm import amberParm
from chemistry import periodic_table

# In the tleap source code:
#
#  bondi    : iGBparm = 0
#  amber6   : iGBparm = 1
#  mbondi   : iGBparm = 2
#  mbondi2  : iGBparm = 6
#  mbondi3  : see below

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

def bondi(parm):
   """ Sets the bondi radii """
   molecule = parm.ToMolecule()
   for i in range(parm.pointers['NATOM']):
      # Radius of C atom depends on what type it is
      if molecule.elements[i] == 'C':
         if parm.parm_data['AMBER_ATOM_TYPE'][i] in ['C1', 'C2', 'C3']:
            parm.parm_data['RADII'][i] = 2.2
         else:
            parm.parm_data['RADII'][i] = 1.7
      # All other elements have fixed radii for all types/partners
      elif molecule.elements[i] == 'H': parm.parm_data['RADII'][i] = 1.2
      elif molecule.elements[i] == 'N': parm.parm_data['RADII'][i] = 1.55
      elif molecule.elements[i] == 'O': parm.parm_data['RADII'][i] = 1.5
      elif molecule.elements[i] == 'F': parm.parm_data['RADII'][i] = 1.5
      elif molecule.elements[i] == 'Si': parm.parm_data['RADII'][i] = 2.1
      elif molecule.elements[i] == 'P': parm.parm_data['RADII'][i] = 1.85
      elif molecule.elements[i] == 'S': parm.parm_data['RADII'][i] = 1.8
      elif molecule.elements[i] == 'Cl': parm.parm_data['RADII'][i] = 1.5
      else: parm.parm_data['RADII'][i] = 1.5

   parm.parm_data['RADIUS_SET'][0] = 'Bondi radii (bondi)'

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

def amber6(parm):
   """ Sets the amber6 radii """
   molecule = parm.ToMolecule()
   for i in range(parm.pointers['NATOM']):
      # Radius of H atom depends on element it is bonded to
      if molecule.elements[i] == 'H':
         if molecule.elements[molecule.bonds[i][0]] in ['C']:
            parm.parm_data['RADII'][i] = 1.3
         elif molecule.elements[molecule.bonds[i][0]] in ['O', 'S']:
            parm.parm_data['RADII'][i] = 0.8
         else:
            parm.parm_data['RADII'][i] = 1.2
      # Radius of C atom depends on what type it is
      elif molecule.elements[i] == 'C':
         if parm.parm_data['AMBER_ATOM_TYPE'][i] in ['C1', 'C2', 'C3']:
            parm.parm_data['RADII'][i] = 2.2
         else:
            parm.parm_data['RADII'][i] = 1.7
      # All other elements have fixed radii for all types/partners
      elif molecule.elements[i] == 'N': parm.parm_data['RADII'][i] = 1.55
      elif molecule.elements[i] == 'O': parm.parm_data['RADII'][i] = 1.5
      elif molecule.elements[i] == 'F': parm.parm_data['RADII'][i] = 1.5
      elif molecule.elements[i] == 'Si': parm.parm_data['RADII'][i] = 2.1
      elif molecule.elements[i] == 'P': parm.parm_data['RADII'][i] = 1.85
      elif molecule.elements[i] == 'S': parm.parm_data['RADII'][i] = 1.8
      elif molecule.elements[i] == 'Cl': parm.parm_data['RADII'][i] = 1.5
      else: parm.parm_data['RADII'][i] = 1.5
   
   parm.parm_data['RADIUS_SET'][0] = 'amber6 modified Bondi radii (amber6)'

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

def mbondi(parm):
   """ Sets the mbondi radii """
   molecule = parm.ToMolecule()
   for i in range(parm.pointers['NATOM']):
      # Radius of H atom depends on element it is bonded to
      if molecule.elements[i] == 'H':
         if molecule.elements[molecule.bonds[i][0]] in ['C', 'N']:
            parm.parm_data['RADII'][i] = 1.3
         elif molecule.elements[molecule.bonds[i][0]] in ['O', 'S']:
            parm.parm_data['RADII'][i] = 0.8
         else:
            parm.parm_data['RADII'][i] = 1.2
      # Radius of C atom depends on what type it is
      elif molecule.elements[i] == 'C':
         if parm.parm_data['AMBER_ATOM_TYPE'][i] in ['C1', 'C2', 'C3']:
            parm.parm_data['RADII'][i] = 2.2
         else:
            parm.parm_data['RADII'][i] = 1.7
      # All other elements have fixed radii for all types/partners
      elif molecule.elements[i] == 'N': parm.parm_data['RADII'][i] = 1.55
      elif molecule.elements[i] == 'O': parm.parm_data['RADII'][i] = 1.5
      elif molecule.elements[i] == 'F': parm.parm_data['RADII'][i] = 1.5
      elif molecule.elements[i] == 'Si': parm.parm_data['RADII'][i] = 2.1
      elif molecule.elements[i] == 'P': parm.parm_data['RADII'][i] = 1.85
      elif molecule.elements[i] == 'S': parm.parm_data['RADII'][i] = 1.8
      elif molecule.elements[i] == 'Cl': parm.parm_data['RADII'][i] = 1.5
      else: parm.parm_data['RADII'][i] = 1.5
   
   parm.parm_data['RADIUS_SET'][0] = 'modified Bondi radii (mbondi)'

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

def mbondi2(parm):
   """ Sets the mbondi2 radii """
   molecule = parm.ToMolecule()
   for i in range(parm.pointers['NATOM']):
      # Radius of H atom depends on element it is bonded to
      if molecule.elements[i] == 'H':
         if molecule.elements[molecule.bonds[i][0]] in ['N']:
            parm.parm_data['RADII'][i] = 1.3
         else:
            parm.parm_data['RADII'][i] = 1.2
      # Radius of C atom depends on what type it is
      elif molecule.elements[i] == 'C':
         if parm.parm_data['AMBER_ATOM_TYPE'][i] in ['C1', 'C2', 'C3']:
            parm.parm_data['RADII'][i] = 2.2
         else:
            parm.parm_data['RADII'][i] = 1.7
      # All other elements have fixed radii for all types/partners
      elif molecule.elements[i] == 'N': parm.parm_data['RADII'][i] = 1.55
      elif molecule.elements[i] == 'O': parm.parm_data['RADII'][i] = 1.5
      elif molecule.elements[i] == 'F': parm.parm_data['RADII'][i] = 1.5
      elif molecule.elements[i] == 'Si': parm.parm_data['RADII'][i] = 2.1
      elif molecule.elements[i] == 'P': parm.parm_data['RADII'][i] = 1.85
      elif molecule.elements[i] == 'S': parm.parm_data['RADII'][i] = 1.8
      elif molecule.elements[i] == 'Cl': parm.parm_data['RADII'][i] = 1.5
      else: parm.parm_data['RADII'][i] = 1.5

   parm.parm_data['RADIUS_SET'][0] = 'H(N)-modified Bondi radii (mbondi2)'

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++

def mbondi3(parm):
   """ Sets mbondi3 radii """
   mbondi2(parm) # start from mbondi2 radii
   molecule = parm.ToMolecule()
   for i in range(parm.pointers['NATOM']):
      # Adjust OE (GLU), OD (ASP), and HH/HE (ARG)
      if parm.parm_data['RESIDUE_LABEL'][molecule.residue_container[i]] in ['GLU','ASP']:
         if parm.parm_data['ATOM_NAME'][i].startswith('OE') or \
            parm.parm_data['ATOM_NAME'][i].startswith('OD'):
            parm.parm_data['RADII'][i] = 1.4
      elif parm.parm_data['RESIDUE_LABEL'][molecule.residue_container[i]] == 'ARG':
         if parm.parm_data['ATOM_NAME'][i].startswith('HH') or \
            parm.parm_data['ATOM_NAME'][i].startswith('HH'):
            parm.parm_data['RADII'][i] = 1.17
      # Adjust carboxylate O radii on C-Termini. Don't just do the end residue,
      # since we can have C-termini in the middle as well (i.e. 2-chain dimers)
      if parm.parm_data['ATOM_NAME'] == 'OXT':
         parm.parm_data['RADII'][i] = 1.4
         parm.parm_data['RADII'][i-1] = 1.4

   parm.parm_data['RADIUS_SET'][0] = 'H(N), salt-bridge modified Bondi radii (mbondi3)'

#++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++


# Begin main part of script now

if __name__ == '__main__':

   # Set up the CL interface

   parser = OptionParser()

   parser.add_option('-p', '--prmtop', dest='prmtop', default='prmtop',
                     help='Change radii of this prmtop (in-place)')
   parser.add_option('-r', '--pbradii', dest='radii', default='mbondi',
                     help='Radius set to change to. May be bondi, amber6, mbondi, mbondi2, or mbondi3')
   parser.add_option('-n', '--new-prmtop', dest='new_prmtop', default=None,
                     help='Name for new topology file. Defaults to overwriting original prmtop')
   
   (opts, args) = parser.parse_args()

   # Process the radii
   opts.radii = opts.radii.lower().strip('"').strip("'")

   # Make sure we've selected valid choices
   if not opts.radii in ['bondi', 'amber6', 'mbondi', 'mbondi2', 'mbondi3']:
      parser.print_help()
      sys.exit(1)

   # Initialize prmtop class and make sure it's OK
   print "Reading in your topology file %s" % opts.prmtop

   parm = amberParm(opts.prmtop)
   if not parm.valid:
      print >> sys.stderr, 'Error: Invalid prmtop (%s)!' % opts.prmtop
      parser.print_help()
      sys.exit(1)

   if opts.new_prmtop == None:
      opts.new_prmtop = opts.prmtop

   print "Changing radius set to %s" % opts.radii

   # Make it so we can overwrite it
   parm.overwrite = True

   # Change the radius set
   if opts.radii == 'bondi':
      bondi(parm)
   elif opts.radii == 'mbondi':
      mbondi(parm)
   elif opts.radii == 'mbondi2':
      mbondi2(parm)
   elif opts.radii == 'mbondi3':
      mbondi3(parm)
   elif opts.radii == 'amber6':
      amber6(parm)
   else:
      print >> sys.stderr, 'Error: Shouldn\'t have gotten here! Bad radius set.'
      sys.exit(1)

   print "Writing out new topology file"

   # Print out the final prmtop
   parm.writeParm(opts.new_prmtop)

   print "Done. Prmtop name %s" % opts.new_prmtop