#!/usr/bin/env python

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

import os, math
import re

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

import Tkinter
from Tkinter import *

import Pmw
import distutils.spawn 
import traceback
import pymol

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

# SimTk modules

import BaseObjectSimTk
import UtilitiesSimTk
import LogSimTk 
import IsimDisplayInterfaceSimTk

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

_ionAbbreviations = None

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

class IsimDisplaySimTk( BaseObjectSimTk.BaseObjectSimTk ): # {

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

   def __init__( self ): # {

      """Constructor for IsimDisplaySimTk"""

      BaseObjectSimTk.BaseObjectSimTk.__init__( self )
 
      # attributes

      self._workingDirectoryName            = None
      self._isimInputFileName               = 'isim.in' 
      self._logReference                    = None 
      self._parameterHash                   = None 
      self._apbsPotentialFileName           = None 
      self._macromoleculeFileName           = None 
      self._isimPotentialFiles              = None 
      self._isimNumberDensityFiles          = None 
      self._isimChargeDensityFiles          = None 
      self._numberOfIonTypes                = -1
      self._temperature                     = -1

      self._requiredDigitsInFileName        = 7     

      # format string used in generating file name

      self._formatString                    = IsimDisplaySimTk.generateFormatString( self._requiredDigitsInFileName )

      # isim.in ion parameters

      self._IonName                         = 'IonName'
      self._IonConcentration                = 'IonConcentration'
      self._IonCharge                       = 'IonCharge'
      self._IonRadius                       = 'IonRadius'
      self._IonLJ_Epsilon                   = 'IonLJ_Epsilon'

      # isim.in parameters

      self._Form                            = 'FORM'
      self._Radius                          = 'RADIUS'
      self._Length                          = 'LENGTH'
      self._Temperature                     = 'TEMPERATURE'
      self._MediumPermittivity              = 'MEDIUM_PERMITTIVITY'
      self._GeneralMode                     = 'GENERAL_MODE'
      self._TotalSteps                      = 'TOTAL_STEPS'
      self._MinimizationSteps               = 'MINIMIZATION_STEPS'
      self._CreationDestructionCycles       = 'CREATION_DESTRUCTION_CYCLES'
      self._ShufflingSteps                  = 'SHUFFLING_STEPS'
      self._ShufflingModes                  = 'SHUFFLING_MODE'
      self._SoManyAreSome                   = 'SO_MANY_ARE_SOME'
      self._MaximumDisplacement             = 'MAXIMUM_DISPLACEMENT'
      self._ParametrizationMinimum          = 'PARAMETRIZATION_MINIMUM'
      self._EquilibrationSteps              = 'EQUILIBRATION_STEPS'
      self._GridResolution                  = 'GRID_RESOLUTION'
      self._CoordinatesOutput               = 'COORDINATES_OUTPUT'
      self._NumberDensitiesOutput           = 'NUMBER_DENSITIES_OUTPUT'
      self._ChargeDensityOutput             = 'CHARGE_DENSITY_OUTPUT'
      self._PotentialOutput                 = 'POTENTIAL_OUTPUT'
      self._PairCorrelationOutput           = 'PAIR_CORRELATION_OUTPUT'
      self._PotentialFrequency              = 'POTENTIAL_FREQUENCY'
      self._PairCorrelationFrequency        = 'PAIR_CORRELATION_FREQUENCY'
      self._PairCorrelationResolution       = 'PAIR_CORRELATION_RESOLUTION'
      self._PairCorrelationShells           = 'PAIR_CORRELATION_SHELLS'
      self._Model                           = 'MODEL'
      self._Types                           = 'TYPES'
      self._TypeId                          = 'TYPE_ID'
      self._StericMode                      = 'STERIC_MODE'
      self._BornIonCharge                   = 'BORN_ION_CHARGE'
      self._BornIonRadius                   = 'BORN_ION_RADIUS'
      self._PqrFileName                     = 'PQR_FILENAME'
      self._PqrDxPotentialFileName          = 'PQR_DX_POTENTIAL_FILENAME'
      self._PqrGridResolution               = 'PQR_GRID_RESOLUTION'
      self._PqrOutsideMeshMode              = 'PQR_OUTSIDE_MESH_MODE'
 
   # } end of __init__

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

   # accessors for ISIM input file 

   def setIsimInputFileName( self, isimInputFileName ): # {
      """Set accessor for IsimInputFileName"""
      self._isimInputFileName = isimInputFileName 
   # } end of setWorkingDirectory

   def getIsimInputFileName( self ): # {
      """Get accessor for IsimInputFileName"""
      return self._isimInputFileName;
   # } end of getIsimInputFileName

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

   # accessors for working directory

   def setWorkingDirectoryName( self, workingDirectoryName ): # {
      """Set accessor for working directory"""
      self._workingDirectoryName = workingDirectoryName
   # } end of setWorkingDirectory

   def getWorkingDirectoryName( self ): # {
      """Get accessor for working directory"""
      return self._workingDirectoryName;
   # } end of getWorkingDirectory

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

   # accessors for Apbs potential file name

   def setApbsPotentialFileName( self, apbsPotentialFileName ): # {
      """Set accessor for ApbsPotentialFileName """
      self._apbsPotentialFileName = apbsPotentialFileName  
   # } end of setApbsPotentialFileName

   def getApbsPotentialFileName( self, removeSuffix = 0 ): # {
      """Get accessor for ApbsPotentialFileName
         Returns name (including suffix)
         Raise ValueError if parameter hash not set or
         if self._TotalSteps not set in parameter hash
      """

      methodName = 'getApbsPotentialFileName: '
      if self._apbsPotentialFileName is None:

         self._apbsPotentialFileName = self.getWorkingDirectoryName()
         if self._parameterHash is None:
             raise ValueError, (methodName + 'Parameter hash is not set -- isim.in file not read?')
         elif self._parameterHash.has_key( self._PqrDxPotentialFileName ):
             self._apbsPotentialFileName = self._parameterHash[self._PqrDxPotentialFileName]
         else:
             raise ValueError, ( methodName + self._PqrDxPotentialFileName + ' not set in parameter hash')

      if removeSuffix:
         return UtilitiesSimTk.UtilitiesSimTk.removeSuffix( self._apbsPotentialFileName );
      else:
         return self._apbsPotentialFileName;

   # } end of getApbsPotentialFileName  

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

   def getApbsPotentialFullPathFileName( self ): # {
      """Get accessor for ApbsPotentialFileName
         Returns full path name
         Raise ValueError if parameter hash not set or
         if self._TotalSteps not set in parameter hash
      """

      methodName = 'getApbsPotentialFullPathFileName: '
      if self._apbsPotentialFileName is None:
         self.getApbsPotentialFileName()
         if self._apbsPotentialFileName is None:
             raise ValueError, (methodName + '_apbsPotentialFileName not set -- isim.in file not read?')

      return self.getWorkingDirectoryName() + self.getFilePathSeparator() + self._apbsPotentialFileName;

   # } end of getApbsPotentialFullPathFileName

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

   # accessors for macromolecule file name

   def setMacromoleculeFileName( self, macromoleculeFileName ): # {
      """Set accessor for macromoleculeFileName """
      self._macromoleculeFileName = macromoleculeFileName  
   # } end of setMacromoleculeFileName

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

   def getMacromoleculeFileName( self, removeSuffix = 0 ): # {
      """Get accessor for macromoleculeFileName
         Returns full path name
         Raise ValueError if parameter hash not set or
         if PqrFileName not set in parameter hash
      """
      methodName = 'getMacromoleculeFileName: '
      if self._macromoleculeFileName is None:
         if self._parameterHash is None:
             raise valueError, methodName + 'Parameter hash is not set -- isim.in file not read?'
         elif self._parameterHash.has_key( self._PqrFileName ):
             self._macromoleculeFileName = self._parameterHash[self._PqrFileName]
         else:
             raise valueError, methodName + self._PqrFileName + ' not available from isim.in file.'

      if removeSuffix:
         return UtilitiesSimTk.UtilitiesSimTk.removeSuffix( self._macromoleculeFileName )
      else:
         return self._macromoleculeFileName

   # } end of getMacromoleculeFileName

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

   def getMacromoleculeFullPathFileName( self ): # {
      """Get accessor for macromoleculeFileName
         Returns full path name
         Raise ValueError if parameter hash not set or
         if PqrFileName not set in parameter hash
      """
      methodName = 'getMacromoleculeFullPathFileName: '
      return self.getWorkingDirectoryName() + self.getFilePathSeparator() + self.getMacromoleculeFileName()

   # } end of getMacromoleculeFullPathFileName

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

   # accessors for ion info

   def getNumberOfIons( self ): # {
      """Get number of ions in simulation
         Raise ValueError if parameter hash not set or
         if 'Types' not set in parameter hash
      """

      methodName = 'getNumberOfIons: '
      if self._numberOfIonTypes < 0:
         if self._parameterHash is None:
             raise valueError, methodName + 'Parameter hash is not set -- isim.in file not read?'
         elif self._parameterHash.has_key( self._Types):
             self._numberOfIonTypes = int( self._parameterHash[self._Types] )
         else:
             raise valueError, methodName + self._Types + ' not available from isim.in file.'

      return self._numberOfIonTypes

   # } end of getNumberOfIons

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

   # accessors for temperature

   def getTemperature( self ): # {
      """Get temperature used in simulation
         Raise ValueError if parameter hash not set or
         if 'Types' not set in parameter hash
      """

      methodName = 'getTemperature: '
      if self._temperature < 0:
         if self._parameterHash is None:
             raise valueError, methodName + 'Parameter hash is not set -- isim.in file not read?'
         elif self._parameterHash.has_key( self._Temperature ):
             self._temperature = float( self._parameterHash[self._Temperature] )
         else:
             self._temperature = 300.0

      return self._temperature

   # } end of getTemperature

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

   def getIonNameGivenIonIndex( self, ionIndex ): # {

      """Get name of ion w/ index=ionIndex
         Raise ValueError if parameter hash not set or
         if key='ionIndex' not set in parameter hash
      """
      return self.getIonFieldGivenIonIndex( ionIndex, self._IonName )

   # } end of getIonNameGivenIonIndex

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

   def getIonChargeGivenIonIndex( self, ionIndex ): # {

      """Get charge of ion w/ index=ionIndex
         Raise ValueError if parameter hash not set or
         if key='ionIndex' not set in parameter hash
      """
      return self.getIonFieldGivenIonIndex( ionIndex, self._IonCharge )

   # } end of getIonChargeGivenIonIndex

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

   def getIonFieldGivenIonIndex( self, ionIndex, ionField ): # {

      """Get ion field for ion w/ index=ionIndex
         Raise ValueError if parameter hash not set or
         if key='ionIndex' not set in parameter hash
      """
      methodName = 'getIonFieldGivenIonIndex: '
      ionDict    = self.getIonDictionaryGivenIonIndex( ionIndex )
      if ionDict.has_key( ionField ):
         return ionDict[ionField]
      else:
         raise valueError, methodName + ' for ionIndex=' + str( ionIndex)  + ' ion ' + ionField + ' is not available from isim.in file.'

   # } end of getIonFieldGivenIonIndex

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

   def getIonDictionaryGivenIonIndex( self, ionIndex ): # {
      """Get name of ion w/ index=ionIndex
         Raise ValueError if parameter hash not set or
         if key='ionIndex' not set in parameter hash
      """

      methodName = 'getIonNameGivenIonIndex: '
      if self._parameterHash is None:
          raise valueError, methodName + 'Parameter hash is not set -- isim.in file not read?'
      elif self._parameterHash.has_key( ionIndex ):
         return self._parameterHash[ionIndex]
      else:
         raise ValueError, methodName + 'ionIndex=' + str( ionIndex)  + ' not available from isim.in file.'

   # } end of getIonDictionaryGivenIonIndex

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

   # build full file name

   def buildFullFileName( self, filePrefix, mcStep, fileSuffix = None, ionIndex = -1 ): # {
      """Build full file name"""

      fullFileName = self.getWorkingPath() + self.buildFileName( filePrefix, mcStep, fileSuffix, ionIndex )

      return fullFileName 

   # } end of buildFullFileName

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

   # build file name

   def buildFileName( self, filePrefix, mcStep, fileSuffix = None, ionIndex = -1 ): # {
      """Build file name"""

      fileName     = self.buildPartialFileNameGivenStep( filePrefix, mcStep, None )

      if ionIndex > 0:
         ionName      = self.getIonNameGivenIonIndex( ionIndex )
         fileName    += '_' + str( ionIndex ) + '_' + ionName 

      if fileSuffix is not None:
         fileName    += fileSuffix

      return fileName 

   # } end of buildFileName

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

   # generate file name give MC step

   def buildPartialFileNameGivenStep( self, filePrefix, step, fileSuffix=None, requiredDigits=None ): # {

      """Helper method to generate file names
      Step is the MC step index
      fileSuffix is optional -- appended to file name
      requiredDigits number of digits in step name (defaults to 7)
      file name prefix is obtained via self.getFilePrefix()
      """

      if requiredDigits is None or requiredDigits == self._requiredDigitsInFileName:
         formatString = self._formatString
      else:   
         formatString = generateFormatString( requiredDigits )

      fileName     = formatString % ( filePrefix, int( step ) )

      if fileSuffix is not None:
         fileName = fileName + fileSuffix

      return fileName

   # } end of buildPartialFileNameGivenStep

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

   # generate format string used in building file names

   def generateFormatString( requiredDigitsInFileName ): # {
      """generateFormatString used in building file names"""
      formatString                          = '%dd' % (requiredDigitsInFileName);
      return                                  '%s%.' + formatString

   generateFormatString=staticmethod( generateFormatString )

   # } end of generateFormatString

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

   # get working path

   def getWorkingPath( self ): # {
      """Get working path"""

      return self.getWorkingDirectoryName() + self.getFilePathSeparator()

   # } end of getWorkingPath

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

   # get ion abbreviations given ion name

   def getIonAbbreviations( ): # {

      """Get singleton dict containing ion abbreviations"""

      global _ionAbbreviations
      if _ionAbbreviations is None:

         _ionAbbreviations               = {}

         _ionAbbreviations['calcium']    = 'Ca'
         _ionAbbreviations['Calcium']    = 'Ca'
         _ionAbbreviations['ca']         = 'Ca'
         _ionAbbreviations['Ca']         = 'Ca'

         _ionAbbreviations['sodium']     = 'Na'
         _ionAbbreviations['Sodium']     = 'Na'
         _ionAbbreviations['na']         = 'Na'
         _ionAbbreviations['Na']         = 'Na'

         _ionAbbreviations['chlorine']   = 'Cl'
         _ionAbbreviations['Chlorine']   = 'Cl'
         _ionAbbreviations['cl']         = 'Cl'
         _ionAbbreviations['Cl']         = 'Cl'

         _ionAbbreviations['magnesium']  = 'Mg'
         _ionAbbreviations['Magnesium']  = 'Mg'
         _ionAbbreviations['mg']         = 'Mg'
         _ionAbbreviations['Mg']         = 'Mg'

         _ionAbbreviations['potassium']  = 'K'
         _ionAbbreviations['Potassium']  = 'K'
         _ionAbbreviations['k']          = 'K'
         _ionAbbreviations['K']          = 'K'

         _ionAbbreviations['ammonium']   = 'NH4'
         _ionAbbreviations['Ammonium']   = 'NH4'
         _ionAbbreviations['nh4']        = 'NH4'
         _ionAbbreviations['Nh4']        = 'NH4'

         _ionAbbreviations['carbonate']  = 'CO3'
         _ionAbbreviations['Carbonate']  = 'CO3'
         _ionAbbreviations['co3']        = 'CO3'
         _ionAbbreviations['CO3']        = 'CO3'

         _ionAbbreviations['nitrate']    = 'NO3'
         _ionAbbreviations['Nitrate']    = 'NO3'
         _ionAbbreviations['no3']        = 'NO3'
         _ionAbbreviations['No3']        = 'NO3'

         _ionAbbreviations['nitrite']    = 'NO2'
         _ionAbbreviations['Nitrite']    = 'NO2'
         _ionAbbreviations['no2']        = 'NO2'
         _ionAbbreviations['NO2']        = 'NO2'

         _ionAbbreviations['phosphate']  = 'PO4'
         _ionAbbreviations['Phosphate']  = 'PO4'
         _ionAbbreviations['po4']        = 'PO4'
         _ionAbbreviations['PO4']        = 'PO4'

         _ionAbbreviations['oxide']      = 'O'
         _ionAbbreviations['Oxide']      = 'O'
         _ionAbbreviations['o']          = 'O'
         _ionAbbreviations['O']          = 'O'

         _ionAbbreviations['acetate']    = 'CH3COO'
         _ionAbbreviations['Acetate']    = 'CH3COO'
         _ionAbbreviations['ch3coo']     = 'CH3COO'
         _ionAbbreviations['CH3COO']     = 'CH3COO'

      return _ionAbbreviations

   getIonAbbreviations=staticmethod( getIonAbbreviations  )

   # } end of getIonAbbreviations 

   def getIonAbbreviationGivenIonName( self, ionName ): # {
      """Get ion abbreviationgiven name
         If ion name is unrecognized, return ion name
         as abbreviation and add { name, name } to dictionary
      """
      
      ionName          = ionName.lower()
      ionAbbreviations = IsimDisplaySimTk.getIonAbbreviations()

      if not ionAbbreviations.has_key( ionName ):
         ionAbbreviations[ionName] = ionName

      return ionAbbreviations[ionName]

   # } end of getIonAbbreviationGivenIonName

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

   # get error message file name

   def getErrorMessageFileName( self ): # {
      """Get error message file name"""
      if self._errorMessageFileName is None:
         self._errorMessageFileName = self.getWorkingDirectoryName() + self.getFilePathSeparator() + 'Error.txt'
      return self._errorMessageFileName

   # } end of getErrorMessageFileName

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

   # generateDisplay

   def generateDisplay( self, app ): # {
      """Generate display"""
      isimDisplayInterface = IsimDisplayInterfaceSimTk.IsimDisplayInterfaceSimTk( app, self )
   # } end of generateDisplay

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

   # read in isim.in file

   def readIsimInputFile( self, isimInputFileName=None ): # {

      """Read ISIM input file; if input isimInputFileName is not set,
         the file = getWorkingDirectory() + '/' + getIsimInputFileName() is used.
         If it is set, then it should be the absolute file path name.
      """

      logReference  = self.getLogReference( )
      methodName    = 'readIsimInputFile:'

      # get Isim input file name, if not provided

      if isimInputFileName == None:
         isimInputFileName = self.getWorkingDirectoryName() + self.getFilePathSeparator() +  self.getIsimInputFileName()

      # open file -- log error if file can not be opened and return None
           
      try:
         isimInputFile = open( isimInputFileName, "r" )
      except IOError:
         logReference.error( methodName + ' file=<' + isimInputFileName + '> could not be opened.' );
         return None

      logReference.info( 'Opened ISIM input file=<' + isimInputFileName + '>' )

      # read input file and load parameters into parameter hash

      isComment           = re.compile( '^#' )
      isIonCount          = re.compile( 'TYPES' )
      self._parameterHash = {}

      while 1:
          nextLine = isimInputFile.readline()
          # print nextLine

          # skip empty lines

          if nextLine == "": break
          nextLine = nextLine[:-1] 

          # skip comments

          if len( nextLine ) < 2 or isComment.match( nextLine ): continue

          # tokenize line and set key => value
          # also check if this is the 'TYPES' entry;
          # if it is, then parse ion name, charge, ...
          # on subsequent lines

          lineTokens                         = nextLine.split()
          self._parameterHash[lineTokens[0]] = lineTokens[1]
          if isIonCount.match( lineTokens[0] ):

             # Sample Ion entry: CALCIUM  5.0  2.0  2.41  0.0

             for ii in range ( 1, (int( lineTokens[1] ) + 1) ):

                nextLine   = isimInputFile.readline()
                nextLine   = nextLine[:-1] 
                lineTokens = nextLine.split()

                self._parameterHash[ii]                           = {}

                self._parameterHash[ii][self._IonName]            = lineTokens[0]
                self._parameterHash[ii][self._IonConcentration]   = float( lineTokens[1] )
                self._parameterHash[ii][self._IonCharge]          = float( lineTokens[2] )
                self._parameterHash[ii][self._IonRadius]          = float( lineTokens[3] )
                self._parameterHash[ii][self._IonLJ_Epsilon]      = float( lineTokens[4] )

      isimInputFile.close()

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

      # echo input file

      numberParameters = len( self._parameterHash );
      logReference.info( 'ISIM input file parameters: ' + str( numberParameters ) )
      tab              = 30
      for key, value in self._parameterHash.iteritems():
         if type( value ) is DictType: 

            for ionKey, ionValue in value.iteritems():
               logReference.info( UtilitiesSimTk.UtilitiesSimTk.getTabbedString( ('Ion' + str( key ) + ' ' + ionKey), (' = ' + str( ionValue )), tab ) )
         else:
            logReference.info( UtilitiesSimTk.UtilitiesSimTk.getTabbedString( key, (' = ' + str( value )), tab ) )

   # } end of readIsimInputFile

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

# } end of class IsimDisplaySimTk