########################################################################
#                                                                      #
# This is a module that contains functions responsible for mutating    #
# the trajectory file for alanine scanning in MMPBSA.py. It must be    #
# included with MMPBSA.py to insure proper functioning of alanine      #
# scanning.                                                            #
#                                                                      #
#          Last updated: 12/18/2009                                    #
#                                                                      #
########################################################################

########################## GPL LICENSE INFO ############################

#  Copyright (C) 2009  Dwight McGee, Billy Miller III, and Jason Swails

#  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 2 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, write to the Free Software
#  Foundation, Inc., 59 Temple Place - Suite 330,
#  Boston, MA 02111-1307, USA.

def MutateMdcrd(traj, newtrajname, prm1, prm2):
   import utils, sys

# number of coordinates per restart line
   wordnum = 10
   
   try:
      mdcrd = open(traj,'r')
   except IOError:
      print >> sys.stderr, 'Error: Error opening ' + traj + '!'
      return -1

   mutres = utils.CheckMutParm(prm1, prm2)

   if mutres == -1:
      return -1

#  Get necessary information from the topology files:
   origresname = utils.getresinfo(mutres, prm1, 'RESIDUE_LABEL')
   number_atoms = utils.numatm(prm1) 
   number_atoms_mut = utils.numatm(prm2)
   resstart = int(utils.getresinfo(mutres, prm1, 'RESIDUE_POINTER'))
   nextresstart = int(utils.getresinfo(mutres + 1, prm1,               \
                      'RESIDUE_POINTER'))

   counter = 0
   frames = 0
   coords_done = 0
   modifying = 0
   coords_tomutate = []
   temp_holder = []
   new_coords = []

# location is 0 before modified coordinates, 1 during modified 
# coordinates, and 2 after modified coordinates
   location = 0   

   if origresname == -1 or resstart == -1:
      print >> sys.stderr, "Error: PRMTOP parsing error! Can't find "  \
             + "RESIDUE_LABEL or RESIDUE_POINTER flag."
      return -1

   if origresname == 'GLY':
      print >> sys.stderr, 'Error: You are trying to mutate GLY to ' + \
               'ALA! Not currently supported.'
      return -1

   new_mdcrd = open(newtrajname, 'w')
# write the new mdcrd here
   for line in mdcrd:

      counter = counter + 1

# first line is always a comment since ptraj always creates this file
      if counter == 1:
         new_mdcrd.write(line.strip() + ' and mutated by MMPBSA.py ' + \
                   'for alanine scanning')
         continue

# Now start reading all coordinates and mutating the relevant ones

# if coordinates of the mutated residue begin on the line we're on...
      if coords_done <= resstart * 3-4 and coords_done + wordnum >=    \
            resstart * 3-4:
         location = 1
         words = line.split()
         for y in range(wordnum):
            words[y] = float(words[y])
         for x in range(wordnum):
            if coords_done <= resstart * 3-4:
               if coords_done % wordnum == 0:
                  new_mdcrd.write('\n')
               new_mdcrd.write(Format(words[x]))
               coords_done = coords_done + 1
            else:
               coords_tomutate.append(words[x])
         continue

      elif location == 1:
         words = line.split()
         for y in range(wordnum):
            words[y] = float(words[y])
         if wordnum + len(coords_tomutate) >= 3 * (nextresstart -      \
                                                            resstart):
            location = 2
            for x in range(wordnum):
               if len(coords_tomutate) < 3 * (nextresstart - resstart):
                  coords_tomutate.append(words[x])
               else:
                  temp_holder.append(words[x])

            new_coords = Mutate(origresname, coords_tomutate)
            for x in range(len(new_coords)):
               if coords_done % wordnum == 0:
                  new_mdcrd.write('\n')
               new_mdcrd.write(Format(new_coords[x]))
               coords_done = coords_done + 1
            if len(temp_holder) != 0:
               for x in range(len(temp_holder)):
                  if coords_done % wordnum == 0:
                     new_mdcrd.write('\n')
                  new_mdcrd.write(Format(temp_holder[x]))
                  coords_done = coords_done + 1
            coords_tomutate = []
            temp_holder = []
         else:
            for x in range(wordnum):
               coords_tomutate.append(words[x])

         continue

      elif location == 2:
         words = line.split()
         for y in range(len(words)):
            words[y] = float(words[y])
         for x in range(len(words)):
            if coords_done % wordnum == 0:
               new_mdcrd.write('\n')
            new_mdcrd.write(Format(words[x]))
            coords_done = coords_done + 1

         if coords_done == number_atoms_mut * 3:
            coords_done = 0
            location = 0
         continue
      else:
         if coords_done % wordnum == 0:
            new_mdcrd.write('\n')
         new_mdcrd.write(line[:len(line)-1])
         coords_done = coords_done + wordnum
         continue

   new_mdcrd.write('\n')

   mdcrd.close()
   new_mdcrd.close()
   return mutres

def Format(string):
   import utils

   return utils.Format(string, 4, 3)

def Mutate(resname, coords):
   import math, utils, sys

   list_one = 'ARG ASH ASN ASP CYM CYS CYX GLH GLN GLU HID HIE HIP ' + \
              'LEU LYN LYS MET PHE SER TRP TYR' 
   list_two = 'ILE THR VAL'
   list_three = 'PRO' 

   chdist = 1.09
   nhdist = 1.01

   coords_tosend = []
   new_coords = []
   coords_received = []
   
   if utils.getnumatms(resname) * 3 == len(coords):
      startindex = 0
      cterm = 0

   elif (utils.getnumatms(resname) + 2) * 3 == len(coords):
      startindex = 2
      cterm = 0

   elif (utils.getnumatms(resname) + 1) * 3 == len(coords):
      startindex = 1
      cterm = 1

   else:
      print >> sys.stderr, utils.getnumatms(resname)
      print >> sys.stderr, len(coords)
      print >> sys.stderr, 'Mismatch in number of atoms in the given ' \
             + 'residue! Check.'
      return -1

   if resname in list_one:

      for x in range((7 + startindex)*3):
         new_coords.append(coords[x])

      for x in range(3):
         coords_tosend.append(coords[(4+startindex)*3+x])
      for x in range(3):
         coords_tosend.append(coords[(7+startindex)*3+x])
      coords_received = utils.ScaleDistance(coords_tosend, chdist)

      for x in range(3):
         new_coords.append(coords_received[x+3])

   elif resname in list_two:

      for x in range((6+startindex)*3):
         new_coords.append(coords[x])

      for x in range(3):
         coords_tosend.append(coords[(4+startindex)*3+x])
      for x in range(3):
         coords_tosend.append(coords[(6+startindex)*3+x])
      coords_received = utils.ScaleDistance(coords_tosend, chdist)

      for x in range(3):
         new_coords.append(coords_received[x+3])

      coords_tosend = []
      coords_received = []
      for x in range(3):
         coords_tosend.append(coords[(4+startindex)*3+x])
      for x in range(3):
         coords_tosend.append(coords[(10+startindex)*3+x])
      coords_received = utils.ScaleDistance(coords_tosend, chdist)

      for x in range(3):
         new_coords.append(coords_received[3+x])

   elif resname in list_three:
      for x in range((1+startindex)*3):
         new_coords.append(coords[x])

      coords_tosend = coords[startindex * 3:startindex * 3 + 6]
      coords_received = utils.ScaleDistance(coords_tosend, nhdist)

      for x in range(3):
         new_coords.append(coords_received[x+3])

      for x in range(6):
         new_coords.append(coords[(10+startindex)*3+x])

      for x in range(9):
         new_coords.append(coords[(7+startindex)*3+x])

      coords_tosend = []
      coords_received = []
      for x in range(3):
         coords_tosend.append(coords[(7+startindex)*3+x])
      for x in range(3):
         coords_tosend.append(coords[(4+startindex)*3+x])
      coords_received = utils.ScaleDistance(coords_tosend, chdist)

      for x in range(3):
         new_coords.append(coords_received[x+3])

   else:
      print >> sys.stderr, 'Error: Residue ' + resname +               \
               ' must have slipped through!'
      return -1

   if cterm == 0:
      for x in range(6):
         new_coords.append(coords[len(coords)-6+x])

   else:
      for x in range(9):
         new_coords.append(coords[len(coords)-9+x])

   return new_coords