"""
This module will create an amber mdin file for either sander or      
pmemd (or others). The program specification loads the appropriate   
dictionaries with default values, etc. It can read and write mdins. 

                           GPL LICENSE INFO                             

Copyright (C) 2009 - 2011 Dwight Mcgee, Bill 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.
"""

# This module will create and read a sander/pmemd input
from MMPBSA_mods.cntrl import cntrl
from MMPBSA_mods.ewald import ewald
from MMPBSA_mods.pb import pb
from MMPBSA_mods.qmmm import qmmm
from MMPBSA_mods.exceptions import CreateInputError
from sys import stderr

#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

def addOn(line, string, file):
   if len(line.strip()) == 0:
      return line + string
   elif len(line) + len(string) > 40:
      file.write(line + '\n')
      return ' ' + string
   else:
      return line + string

#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+

class mdin:
   
   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def __init__(self, program = 'sander', verbosity = 1):
      # define instance data
      self.program = program   # which program we're creating the input file for
      self.cntrl_obj = cntrl() # object with cntrl namelist vars in a dictionary
      self.ewald_obj = ewald() # object with ewald namelist vars in a dictionary
      self.pb_obj = pb()       # object with pb namelist vars in a dictionary
      self.qmmm_obj = qmmm()   # object with qmmm namelist vars in a dictionary
      self.verbosity = 0       # verbosity level: 0 -- print nothing
                               #                  1 -- print errors
                               #                  2 -- 1 + warnings
                               #                  3 -- 2 + notes
      self.cards = []          # array that has all of the input cards that come
                               # after namelists
      self.cntrl_nml = {}      # dictionary with cntrl namelist vars
      self.cntrl_nml_defaults = {} # dictionary with default cntrl namelist vars
      self.ewald_nml = {}          # dictionary with ewald namelist vars
      self.ewald_nml_defaults = {} # dictionary with default ewald namelist vars
      self.pb_nml = {}             # dictionary with pb namelist vars
      self.pb_nml_defaults = {}    # dictionary with default pb namelist vars
      self.qmmm_nml = {}           # dictionary with qmmm namelist vars
      self.qmmm_nml_defaults = {}  # dictionary with default qmmm namelist vars
      self.valid_namelists = []    # array with valid namelists for each program
      self.title = 'mdin prepared by mdin.py'   # title for the mdin file


      if self.program == "sander":
         self.cntrl_nml = self.cntrl_obj.sander
         self.ewald_nml = self.ewald_obj.sander
         self.pb_nml = self.pb_obj.sander
         self.qmmm_nml = self.qmmm_obj.sander
         self.valid_namelists = ['cntrl','ewald','qmmm','pb']
      elif self.program == "sander.APBS":
         self.cntrl_nml = self.cntrl_obj.sander
         self.pb_nml = self.pb_obj.sanderAPBS
         self.valid_namelists = ['cntrl','apbs']
      elif self.program == "pmemd":
         self.cntrl_nml = self.cntrl_obj.pmemd
         self.ewald_nml = self.ewald_obj.pmemd
         self.valid_namelists = ['cntrl','ewald']
      else:
         raise CreateInputError('Error: program (%s) unrecognized!' % self.program)

      self.cntrl_nml_defaults = self.cntrl_nml.copy()
      self.ewald_nml_defaults = self.ewald_nml.copy()
      self.pb_nml_defaults = self.pb_nml.copy()
      self.qmmm_nml_defaults = self.qmmm_nml.copy()

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def write(self, filename = 'mdin'):
      
      if not self.check(): # check the validity of the parameters
         raise CreateInputError('mdin file failed sanity check')

      # open the file for writing and write the header and &cntrl namelist
      file = open(filename,'w')
      file.write(self.title + '\n')
      file.write('&cntrl\n')
      # automatic indent of single space
      line = ' '
      # add any variable that is different from the default to the mdin file
      for var in self.cntrl_nml.keys():
         if self.cntrl_nml[var] != self.cntrl_nml_defaults[var]:
            line = addOn(line, "%s=%s, " % (var, self.cntrl_nml[var]), file)

      # flush any remaining items that haven't yet been printed to the mdin file
      if len(line.strip()) != 0:
         file.write(line + '\n')

      # end the namelist
      file.write('/\n')

      # print the ewald namelist if any variables differ from the default
      line = ' '
      has_been_printed = False  # keep track if this namelist has been printed
      for var in self.ewald_nml.keys():
         if self.ewald_nml[var] != self.ewald_nml_defaults[var]:
            if (not has_been_printed):
               file.write('&ewald\n')
               has_been_printed = True
            line = addOn(line, "%s=%s, " % (var, self.ewald_nml[var]), file)

      # flush any remaining items that haven't been printed to the mdin file
      if len(line.strip()) != 0:
         file.write(line + '\n')

      # end the namelist
      if has_been_printed:
         file.write('/\n')

      # print the pb namelist if any variables differ from the original
      line = ' '
      has_been_printed = False # keep track if this namelist has been printed
      for var in self.pb_nml.keys():
         if self.pb_nml[var] != self.pb_nml_defaults[var]:
            if (not has_been_printed):
               if self.program == 'sander.APBS':
                  file.write('&apbs\n')
               else:
                  file.write('&pb\n')
               has_been_printed = True
            line = addOn(line,'%s=%s, ' % (var, self.pb_nml[var]), file)

      # flush any remaining items that haven't been printed to the mdin file
      if len(line.strip()) != 0:
         file.write(line + '\n')

      # end the namelist
      if has_been_printed:
         file.write('/\n')

      # print the qmmm namelist if any variables differ from the original
      line = ' '
      has_been_printed = False # keep track if this namelist has been printed
      if self.cntrl_nml['ifqnt'] == 1:
         for var in self.qmmm_nml.keys():
            if self.qmmm_nml[var] != self.qmmm_nml_defaults[var]:
               if (not has_been_printed):
                  file.write('&qmmm\n')
                  has_been_printed = True
               line = addOn(line, '%s=%s, ' % (var, self.qmmm_nml[var]), file)
      
         # flush any remaining items that haven't been printed to the mdin file
         if len(line.strip()) != 0:
            file.write(line + '\n')

         # end the namelist
         if has_been_printed:
            file.write('/\n')

      # Write the cards to the input file
      for i in range(len(self.cards)):
         file.write(self.cards[i].strip() + '\n')
      if len(self.cards) != 0:
         file.write('END\n')
      
      file.close()

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#
   
   def read(self, filename = 'mdin'):
      try:
         file = open(filename, 'r')
      except IOError:
         print >> stderr, 'File (%s) can\'t be opened for reading...' % filename
         return -1

      lines = file.readlines()
      file.close()

      # split up input file into separate fields by comma
      blocks = []  # namelists in the order they appear
      block_fields = [] # array of arrays that correspond to entries in 
                        # namelists found in "blocks" above
      inblock = False
      lead_comment = True
      for i in range(len(lines)):
         if not inblock and not lines[i].strip().startswith('&') and \
                lead_comment:
            continue
         elif not inblock and not lines[i].strip().startswith('&') and \
                not lead_comment:
            final_ended = True
            for j in range(i,len(lines)):
               if lines[j].strip().startswith('&'):
                  final_ended = False
            if final_ended and len(lines[i].strip()) != 0:
               self.cards.append(lines[i])
         elif not inblock and lines[i].strip().startswith('&'):
            lead_comment = False
            inblock = True
            block = lines[i].strip()[1:].lower()
            blocks.append(block)    # add the name of the namelist to "blocks"
            block_fields.append([]) # add empty array to be filled with entries 
                                    # for given namelist
            if not block in self.valid_namelists:
               raise CreateInputError('Invalid namelist (%s) in input file (%s) for %s'
                             % (lines[i].strip(), filename, self.program))
         elif inblock and (lines[i].strip() == '/' or 
                           lines[i].strip() == '&end'):
            inblock = False
         elif inblock and lines[i].strip().startswith('&'):
            raise CreateInputError(('Invalid input file (%s). Terminate each ' +
                             'namelist before another is started') % filename)
         elif inblock:
            items = lines[i].strip().split(',')
            j = 0
            while j < len(items):
               items[j] = items[j].strip()
               if len(items[j]) == 0:
                  items.pop(j)
               else:
                  j += 1
            block_fields[len(block_fields)-1].extend(items)

      # take out the last END in the cards if it's there
      if len(self.cards) != 0 and \
            self.cards[len(self.cards)-1].strip().upper() == 'END':
         self.cards.pop()

      # combine any multi-element fields: e.g. rstwt=1,2,3,
      begin_field = -1
      for i in range(len(block_fields)):
         for j in range(len(block_fields[i])):
            if not '=' in block_fields[i][j]:
               if begin_field == -1:
                  raise CreateInputError('Invalid input file (%s).' % filename)
               else:
                  block_fields[i][begin_field] += ',' + block_fields[i][j]
            else:
               begin_field = j

      # now parse through the options and add them to the dictionaries
      for i in range(len(block_fields)):
         for j in range(len(block_fields[i])):
            if not '=' in block_fields[i][j]:
               continue
            else:
               var = block_fields[i][j].split('=')
               self.change(blocks[i], var[0].strip(), var[1].strip())

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def change(self, namelist, variable, value):
      """ Change the value of a variable without adding a new key-pair """
      
      variable = variable.lower()
      
      if namelist == "cntrl":
         if variable in self.cntrl_nml.keys():
            self.cntrl_nml[variable] = value
         else:
            raise CreateInputError('Unknown variable (%s) in &cntrl!' % variable)
      elif namelist == 'ewald': 
         if variable in self.ewald_nml.keys():
            self.ewald_nml[variable] = value
         else:
            raise CreateInputError('Unknown variable (%s) in &ewald!' % variable)
      elif namelist == 'pb' or namelist == 'apbs':
         if variable in self.pb_nml.keys():
            self.pb_nml[variable] = value
         else:
            raise CreateInputError('Unknown variable (%s) in &%s!' % 
                            (variable, namelist))
      elif namelist == 'qmmm':
         if variable in self.qmmm_nml.keys():
            self.qmmm_nml[variable] = value
         else:
            raise CreateInputError('Unknown variable (%s) in &qmmm' % variable)
      else:
         raise CreateInputError('Unknown namelist (%s)!' % namelist)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def check(self):
      return True

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def SHAKE(self):
      self.change('cntrl','ntf', 2)
      self.change('cntrl','ntc', 2)
      self.change('cntrl','dt', 0.002)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def constPressure(self, press=1.0, taup=1.0):
      self.change('cntrl','ntb', 2)
      self.change('cntrl','ntp', 1)
      self.change('cntrl','pres0', press)
      self.change('cntrl','taup', taup)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def constVolume(self):
      self.change('cntrl','ntb', 1)
      self.change('cntrl','ntp', 0)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def constTemp(self, ntt=3, temp=300.0, gamma_ln=2.0, ig=-1, tautp=1.0):
      self.change('cntrl','ntt', ntt)
      self.change('cntrl','temp0', temp)
      self.change('cntrl','tempi', temp)
      self.change('cntrl','gamma_ln', gamma_ln)
      self.change('cntrl','ig', ig)
      self.change('cntrl','tautp', tautp)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def constpH(self, solvph=7.0, igb=2, ntcnstph=10):
      self.change('cntrl','icnstph', 1)
      self.change('cntrl','solvph', solvph)
      self.change('cntrl','ntcnstph', ntcnstph)
      self.change('cntrl','igb', igb)
      self.change('cntrl','ntb', 0)
      self.change('cntrl','saltcon', 0.1)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def restrainHeavyAtoms(self, restraint_wt=0.0):
      self.change('cntrl','ntr', 1)
      self.change('cntrl','restraint_wt', restraint_wt)
      self.change('cntrl','restraintmask', '!@H=')

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def restrainBackbone(self, restraint_wt=0.0):
      self.change('cntrl','ntr', 1)
      self.change('cntrl','restraint_wt', restraint_wt)
      self.change('cntrl','restraintmask', '@N,CA,C')

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def genBorn(self, igb=5, rgbmax=25.0):
      self.change('cntrl','igb', igb)
      self.change('cntrl','ntb', 0)
      self.change('cntrl','ntp', 0)
      self.change('cntrl','rgbmax', rgbmax)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def time(self, time=1000.0, dt=-1): # time in ps
      if dt == -1:
         if self.cntrl_nml['ntc'] == 2 and self.cntrl_nml['ntf'] == 2:
            dt = 0.002
         else:
            dt = 0.001
      time = int(time / dt)

      self.change('cntrl','dt', dt)
      self.change('cntrl','nstlim', time)
      self.change('cntrl','imin', 0)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def heat(self, tempi=0.0, temp0=300.0, ntt=3, tautp=5.0, ig=-1, gamma_ln=5.0):
      self.constVolume()
      self.change('cntrl','tempi', tempi)
      self.change('cntrl','temp0', temp0)
      self.change('cntrl','ntt', ntt)
      self.change('cntrl','tautp', tautp)
      self.change('cntrl','ig', ig)
      self.change('cntrl','gamma_ln', gamma_ln)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def restart(self,ntx=5):
      self.change('cntrl','irest',1)
      self.change('cntrl','ntx',ntx)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def TI(self, clambda=0.0):
      self.change('cntrl','clambda', clambda)
      self.change('cntrl','icfe',1)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def softcore_TI(self, scalpha=0.5, scmask='', crgmask='', logdvdl=0):
      self.change('cntrl','icfe',1)
      self.change('cntrl','ifsc',1)
      self.change('cntrl','scalpha',scalpha)
      self.change('cntrl','scmask',scmask)
      self.change('cntrl','crgmask',crgmask)
      self.change('cntrl','logdvdl',logdvdl)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def minimization(self, imin=1, maxcyc=1, ncyc=10, ntmin=1):
      self.change('cntrl','imin', imin)
      self.change('cntrl','maxcyc', maxcyc)
      self.change('cntrl','ncyc', ncyc)
      self.change('cntrl','ntmin', ntmin)

   #-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#-#

   def AddCard(self, title='Residues in card', cardString='RES 1'):
      self.cards.append('%s\n%s\nEND' % (title, cardString))

#+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+