
# this is quick and dirty irst crack at a gamess interface, written by someone who
# knows very little about the program (WLD) - hence the version 1 identifier...

# by the way, most of the below is untested...
      
import os
import shutil
import glob
import re
import string
import sys
import time

from chempy import feedback
from chempy.brick import Brick

atNum = {
   'H'  : 1,
   'C'  : 6,
   'N'  : 7,
   'O'  : 8,
   'F'  : 9,
   'P'  : 15,
   'S'  : 16,
   'Cl' : 17,
   'Br' : 35,
   'I'  : 53,
   }

def do(input,run_prefix=None,echo=None,
       punch=None,output=None,skip=None):
   if not run_prefix:
      run_prefix = 'gamess_run'
   if not skip:
      if feedback['gamess']:
         print " "+str(__name__)+': creating temporary files "%s.*"' % (run_prefix)
         print " "+str(__name__)+': launching gamess...' 
      try:
         for a in glob.glob(run_prefix+".*"):
            os.unlink(a)
      except:
         pass
      f = open(run_prefix+".inp",'w')
      for a in input:
         f.write(a)
      f.close()
      if echo:
         os.system(rungms_path+' '+run_prefix+" 2>&1 | tee "+run_prefix+".out")
      else:
         os.system(rungms_path+' '+run_prefix+" > "+run_prefix+".out 2>&1")      
# NFS workaround (flushes the directory cache so that glob will work)
      try: os.unlink(".sync")
      except: pass
      f = open(".sync",'w')
      f.close()
#
   if feedback['gamess']:
      print " "+str(__name__)+': job complete. '
   if punch:
      for src in glob.glob(run_prefix+".dat"):
         f = open(src)
         punch = f.readlines()
         f.close()
   if output:
      for src in glob.glob(run_prefix+".out"):
         f = open(src)
         output = f.readlines()
         f.close()
   return (output,punch)

if os.environ.has_key('GAMESS'):
   base = os.environ['GAMESS']
   bin_path = base + '/bin/'
   rungms_path = bin_path + 'rungms'
else:
   base = ''
   bin_path = ''
   params_path = ''

class State:

   def __init__(self):
      self.model = None
      self.data = None
      self.vec = None
      
   def load_model(self,model):
      self.model = model

   def get_zmat_ordering(self):
      lst = []
      for z in self.model.get_internal_tuples():
         lst.append(z[0])
      return lst
   
   def get_data_group(self,basis = None,zmat = 1):
      model = self.model
         
      gmsList = []
      
      # write header records
      gmsList.append(" $DATA\n")
      gmsList.append(model.molecule.title+" from "+str(__name__)+"\n")
      gmsList.append("C1\n")

      # write atom records in an ordering compatible with internal
      # coordinate generation
      c = 1
      for z in self.get_zmat_ordering():
         a = model.atom[z]
         if not len(a.name):
            name = a.symbol + "%02d"%c
         else:
            name = a.name
         gmsList.append("%10s %5.1f %18.10f %18.10f %18.10f\n" %
                        (name,atNum[a.symbol],a.coord[0],
                         a.coord[1],a.coord[2]))
         c = c + 1
      gmsList.append(" $END\n")
      return gmsList

   def get_ordered_data_group(self): 
      gmsList = self.data[0:3]
      flag = 1
      c = 3
      for a in self.data[3:]:
         if flag:
            flag = 0
            gmsList.append(a)
         if string.strip(a)=='':
            flag = 1
         c = c + 1
      return gmsList
   
   def get_contrl_group(self,
                        scftyp='RHF',
                        runtyp='ENERGY',
                        exetyp='RUN',
                        coord='UNIQUE',
                        nzvar = -1):
      gmsList = []
      model = self.model
      if nzvar:
         if nzvar<0:
            nzvar = (self.model.nAtom*3)-6
      gmsList.append(" $CONTRL SCFTYP=%s RUNTYP=%s EXETYP=%s\n"
                     % (scftyp,runtyp,exetyp) )
      if coord:
         gmsList.append("COORD=%s\n"%coord)
      if nzvar:
         gmsList.append("NZVAR=%d\n"%nzvar)
      chg = 0
      for a in model.atom:
         chg = chg + a.formal_charge
      chg = int(chg)
      if chg==0:
         icharg=None
      else:
         icharg='ICHARG=%d' % chg
      if icharg:
         gmsList.append("%s\n" % (icharg))
      gmsList.append(" $END\n")
      return gmsList

   def read_output_list(self,list):
      ll = len(list)
      c = 0
      crd_list = []
      chg_list = []
      nrg_list = []
      for a in list:
         if a[0:36] == ' COORDINATES OF ALL ATOMS ARE (ANGS)':
            crd_list.append(c+3)
         if a[0:13] == ' NET CHARGES:':
            chg_list.append(c+4)
         if a[0:37] == '                       TOTAL ENERGY =':
            nrg_list.append(c)
         c = c + 1
      atom = self.model.atom
      idx = {}
      c = 0
      for a in atom:
         idx[string.upper(a.name)]=c  # games converts to uppercase
         c = c + 1
      if len(crd_list):
         a = crd_list.pop()
         cc = 0
         while a<ll:
            l = list[a]
            name = string.strip(l[1:11])
            if name=='':
               break
            atom[idx[name]].coord = [float(l[16:31]),
                                     float(l[31:46]),
                                     float(l[46:61])]
            cc = cc + 1
            a = a + 1
         if cc and feedback['gamess']:
            print " "+str(__name__)+': coordinates modified for %d atoms.' % (cc)
      if len(chg_list):
         a = chg_list.pop()
         cc = 0
         while a<ll:
            l = list[a]
            name = string.strip(l[1:11])
            if name[0]=='-':
               break
            atom[idx[name]].partial_charge = float(l[19:27])
            a = a + 1
            cc = cc + 1
         if cc and feedback['gamess']:
            print " "+str(__name__)+': charges modified for %d atoms.' % (cc)
      if len(nrg_list):
         a = nrg_list.pop()
         l = list[a]
         # get energy, and convert to kcal/mole
         self.model.molecule.energy = float(string.strip(l[38:58]))*627.5095
         if feedback['gamess']:
            print " "+str(__name__)+': energy updated %12.6f.' % self.model.molecule.energy
         
   def read_punch_list(self,list):
      ll = len(list)
      c = 0
      data_list = []
      vec_list = []
      for a in list:
         if a[0:6] == ' $DATA':
            data_list.append(c)
         elif a[0:5] == ' $VEC':
            vec_list.append(c)
         c = c + 1
      if len(data_list):
         a = data_list.pop()
         self.data = []
         data = self.data
         while a<ll:
            la = list[a]
            data.append(la)
            if la[0:5] == ' $END':
               break
            a = a + 1
         if feedback['gamess']:
            print " "+str(__name__)+': read $DATA group.'
      if len(vec_list):
         a = vec_list.pop()
         self.vec = []
         vec = self.data
         while a<ll:
            la = list[a]
            vec.append(la)
            if la[0:5] == ' $END':
               break
            a = a + 1
         if feedback['gamess']:
            print " "+str(__name__)+': read new $VEC group.'

   def update_data_coords(self): # update coordinates of ordered atoms in $DATA
      idx = {}
      c = 0
      for a in self.model.atom:
         idx[string.upper(a.name)]=c
         c = c + 1
      if self.data:
         flag = 1
         c = 3
         for a in self.data[3:]:
            if flag:
               flag = 0
               kee = a[0:3]
               if not idx.has_key(kee):
                  break
               i = idx[kee]
               at = self.model.atom[i]
               self.data[c]="%-10s%5.1f%18.10f%18.10f%18.10f\n" % (
                  at.name,atNum[at.symbol],at.coord[0],
                  at.coord[1],at.coord[2])
            if string.strip(a)=='':
               flag = 1
            c = c + 1
   
   def read_density_list(self,list,brick,z_step):
      ll = len(list)
      c = 0
      den_list = []
      for a in list:
         if a[0:37] == ' ELECTRON DENSITY, IPOINT,X,Y,Z,EDENS':
            den_list.append(c+1)
         c = c + 1
      if len(den_list):
         lst = 0
         a = den_list.pop()
         for x in xrange(brick.dim[0]):
            for y in xrange(brick.dim[1]):
               brick.lvl[x][y][z_step] = float(list[a][36:51])
               a = a + 1 
         if feedback['gamess']:
            print " "+str(__name__)+': read density slice %d of %d.' %(
               z_step+1,brick.dim[2])

   def read_potential_list(self,list,brick,z_step):
      ll = len(list)
      c = 0
      pot_list = []
      for a in list:
         if a[0:51] == 'THE ROWS OF THE ELECTROSTATIC POTENTIAL GRID (A.U.)':
            pot_list.append(c+1)
         c = c + 1
      if len(pot_list):
         lst = 0
         a = pot_list.pop()
         for x in xrange(brick.dim[0]):
            aa = a
            mat = list[aa][0:3]
            col = []
            while 1:
               if list[aa][0:3]==mat:
                  col.append(list[aa][5:])
                  aa = aa + 1
               else:
                  break
            a = aa
            vst = string.split(string.strip(string.join(col)))
            for y in xrange(brick.dim[1]):
               brick.lvl[x][y][z_step] = float(vst[y])
         if feedback['gamess']:
            print " "+str(__name__)+': read potential slice %d of %d.' %(
               z_step+1,brick.dim[2])

   def get_basis_group(self,gbasis='N31',ngauss=6,ndfunc=1):
      gmsList = []
      gmsList.append(" $BASIS GBASIS=%s NGAUSS=%d NDFUNC=%d\n" %
                     (gbasis,ngauss,ndfunc))
      model = self.model
      chg = 0
      for a in model.atom:
         chg = chg + a.formal_charge
      chg = int(chg)
      if chg<0:
         diffsp=' DIFFSP=.TRUE.'
      else:
         diffsp=None
      if diffsp:
         gmsList.append("%s\n" % (diffsp))
      gmsList.append(" $END\n")
      return gmsList

   def get_zmat_ext_freeze_torsion(self,flag=3):
      # requires PYMOL to read dihedrals from structure
      # requires list of dihedrals from tinker.amber
      #
      from pymol import cmd
      from tinker.amber import Topology      

      cmd.load_model(self.model,'_gamess1')
      model = self.model

      # get mapping of model ordering to zmat ordering
      m2z = {}
      z2m = {}
      c = 1 # GAMESS is one-based
      for a in self.get_zmat_ordering():
         m2z[a] = c
         z2m[c] = a
         c = c + 1

      # get all torsions in the molecule

      topo = Topology(self.model)

      # find those where flag is set in all atoms

      mask = 2 ** flag

      frozen_list = []

      for a in topo.torsion.keys():
         if (model.atom[a[0]].flags&
             model.atom[a[1]].flags&
             model.atom[a[2]].flags&
             model.atom[a[3]].flags)&mask:
            frozen_list.append(a)

      print " freeze-torsion: %d torsions will be frozen."%len(frozen_list)

      irzmat = []
      ifzmat = []
      fvalue = []
      if len(frozen_list):

         for frozen in frozen_list:
            # find additional torsions which need to be removed

            remove = []

            for a in topo.torsion.keys():
               if (((a[1]==frozen[1])and(a[2]==frozen[2])) or
                   ((a[2]==frozen[1])and(a[1]==frozen[2]))):
                  if a!=frozen:
                     remove.append(a)

            # convert to internal coordinate ordering

            frozen_z = (m2z[frozen[0]],m2z[frozen[1]],
                        m2z[frozen[2]],m2z[frozen[3]])

            remove_z = []
            for a in remove:
               remove_z.append(m2z[a[0]],m2z[a[1]],m2z[a[2]],m2z[a[3]])

            # now reorder atoms in torsions to reflect z_matrix ordering
            # (not sure this is necessary)

            if frozen_z[0]>frozen_z[3]:
               frozen_z = (frozen_z[3],frozen_z[2],frozen_z[1],frozen_z[0])
            tmp_z = []
            for a in remove_z:
               if a[0]>a[3]:
                  tmp_z.append((a[3],a[2],a[1],a[0]))
               else:
                  tmp_z.append(a)
            remove_z = tmp_z

            # get value of the fixed torsion

            fixed = (z2m[frozen_z[0]],z2m[frozen_z[1]],
                     z2m[frozen_z[2]],z2m[frozen_z[3]])

            dihe = cmd.get_dihedral("(_gamess1 and id %d)"%fixed[0],
                             "(_gamess1 and id %d)"%fixed[1],
                             "(_gamess1 and id %d)"%fixed[2],
                             "(_gamess1 and id %d)"%fixed[3])

            # write out report for user edification

            print " freeze-torsion: fixing freeze-torsion:"
            print " freeze-torsion: %d-%d-%d-%d (pymol), %d-%d-%d-%d (gamess)"%(
               fixed[0],fixed[1],fixed[2],fixed[3],
               frozen_z[0],frozen_z[1],frozen_z[2],frozen_z[3])
            print " freeze-torsion: at %5.3f"%dihe
            print " freeze-torsion: removing redundant torsions:"
            for a in remove_z[1:]:
               print " freeze-torsion: %d-%d-%d-%d (pymol), %d-%d-%d-%d (gamess)"%(
                  z2m[a[0]],z2m[a[1]],z2m[a[2]],z2m[a[3]],
                  a[0],a[1],a[2],a[3])

            # add parameters for this torsion into the list

            ifzmat.append((3,frozen_z[0],frozen_z[1],frozen_z[2],frozen_z[3]))
            fvalue.append(dihe)

            if len(remove_z):
               for a in remove_z[1:]:
                  irzmat.append((3,a[0],a[1],a[2],a[3]))

      # generate restrained dihedral information

      zmat_ext = []
      if len(ifzmat):
         zmat_ext.append(" IFZMAT(1)=\n")
         comma = ""
         for a in ifzmat:
            zmat_ext.append(comma+"%d,%d,%d,%d,%d\n"%a)
            comma = ","
      if len(fvalue):
         zmat_ext.append(" FVALUE(1)=\n")
         comma = ""
         for a in fvalue:
            zmat_ext.append(comma+"%1.7f\n"%a)
            comma = ","
      if len(irzmat):
         zmat_ext.append(" IRZMAT(1)=\n")
         comma = ""
         for a in irzmat:
            zmat_ext.append(comma+"%d,%d,%d,%d,%d\n"%a)
            comma = ","

      cmd.delete("_gamess1") # important
      if len(zmat_ext):
         return zmat_ext
      else:
         return None
      
   def get_zmat_group(self,auto=1,dlc=1,zmat_extend=None):
      gmsList = []
      if auto and dlc:
         if zmat_extend == None:
            gmsList.append(" $ZMAT DLC=.TRUE. AUTO=.TRUE. $END\n")
         else:
            gmsList.append(" $ZMAT DLC=.TRUE. AUTO=.TRUE.\n")
            gmsList.extend(zmat_extend)
            gmsList.append(" $END\n")            
      else:
         raise RuntimeError
      return gmsList
   
   def get_eldens_group(self,morb=0):
      gmsList = []
      gmsList.append(" $ELDENS IEDEN=1 MORB=%i \n" % morb)
      gmsList.append("WHERE=GRID OUTPUT=PUNCH\n $END\n")
      return gmsList

   def get_elpot_group(self,morb=0):
      gmsList = []
      gmsList.append(" $ELPOT IEPOT=1 \n")
      gmsList.append("WHERE=GRID OUTPUT=PUNCH\n $END\n")
      return gmsList

   def get_guess_group(self,guess='HUCKEL'):
      return [" $GUESS GUESS=%s $END\n"%guess]
   
   def get_grid_group(self,brick,z_step):
      origin = (
         brick.origin[0],
         brick.origin[1],
         brick.origin[2]+brick.grid[2]*z_step)
      x_coord = (
         brick.origin[0]+brick.range[0]+brick.grid[0]/100.0,
         brick.origin[1],
         brick.origin[2]+brick.grid[2]*z_step)
      y_coord = (
         brick.origin[0],
         brick.origin[1]+brick.range[1]+brick.grid[1]/100.0,
         brick.origin[2]+brick.grid[2]*z_step)
      gmsList = [
         " $GRID ORIGIN(1)=%12.5f,%12.5f,%12.5f\n" % origin,
         "XVEC(1) = %12.5f,%12.5f,%12.5f\n" % x_coord,
         "YVEC(1) = %12.5f,%12.5f,%12.5f\n" % y_coord,
         "SIZE = %12.5f\n" % brick.grid[0],
         " $END\n"
         ]
      return gmsList
   
   def get_scf(self,dirscf=1):
      gmsList = []
      if dirscf:
         gmsList.append(" $SCF DIRSCF=.TRUE. $END\n")
      return gmsList

   def get_optimize_job(self,dirscf=1,zmat_extend=None):
      gmsList = []
      gmsList.extend(self.get_contrl_group(runtyp='OPTIMIZE'))
      gmsList.extend(self.get_basis_group())
      gmsList.extend(self.get_scf(dirscf=dirscf))
      gmsList.extend(self.get_data_group())
      gmsList.extend(self.get_zmat_group(zmat_extend=zmat_extend))
      gmsList.append(" $STATPT NSTEP=50 $END\n")
      return gmsList
         
   def get_optimize_charge_job(self):
      gmsList = self.get_optimize_job()
      gmsList.append(" $ELPOT IEPOT=1 WHERE=PDC $END\n")
      gmsList.append(" $PDC PTSEL=GEODESIC CONSTR=CHARGE $END\n")
      return gmsList

   def get_energy_charge_job(self):
      gmsList = self.get_energy_job()
      gmsList.append(" $ELPOT IEPOT=1 WHERE=PDC $END\n")
      gmsList.append(" $PDC PTSEL=GEODESIC CONSTR=CHARGE $END\n")
      return gmsList
   
   def get_prop_job(self):
      gmsList = []
      gmsList.extend(self.get_contrl_group(runtyp = 'PROP',
                                           nzvar=0))
      gmsList.extend(self.get_guess_group(guess='MOREAD'))
      self.update_data_coords()
      gmsList.extend(self.data)
      gmsList.extend(self.vec)
      return gmsList
      
   def get_energy_job(self):
      gmsList=[]
      gmsList.extend(self.get_contrl_group(
         runtyp = 'ENERGY'
         ))
      gmsList.extend(self.get_basis_group())
      gmsList.extend(self.get_scf())
      gmsList.extend(self.get_data_group())
      gmsList.extend(self.get_zmat_group())
      return gmsList

   def get_density_job(self,brick,z_step,morb=0):
      gmsList = self.get_prop_job()
      gmsList.extend(self.get_eldens_group(morb=morb))
      gmsList.extend(self.get_grid_group(brick,z_step))
      return gmsList

   def get_potential_job(self,brick,z_step):
      gmsList = self.get_prop_job()
      gmsList.extend(self.get_elpot_group())
      gmsList.extend(self.get_grid_group(brick,z_step))
      return gmsList

   def get_prop_charge_job(self):
      gmsList = self.get_prop_job()
      gmsList.append(" $ELPOT IEPOT=1 WHERE=PDC $END\n")
      gmsList.append(" $PDC PTSEL=GEODESIC CONSTR=CHARGE $END\n")
      return gmsList

   def get_density(self,brick,morb=0,run_prefix=None):
      for a in xrange(brick.dim[2]):
         gmsList = self.get_density_job(brick,a,morb=morb)
         result = do(gmsList,punch=1,
                     run_prefix=run_prefix)
         self.read_density_list(result[1],brick,a)

   def get_potential(self,brick,run_prefix=None):
      for a in xrange(brick.dim[2]):
         gmsList = self.get_potential_job(brick,a)
         result = do(gmsList,punch=1,
                     run_prefix=run_prefix)
         self.read_potential_list(result[1],brick,a)

   def get_charges(self,run_prefix=None):
      gmsList = self.get_energy_job()
      result = do(gmsList,output=1,punch=1,
                  run_prefix=run_prefix)
      self.read_output_list(result[0])
      self.read_punch_list(result[1])

   def get_energy(self,run_prefix=None):
      gmsList = self.get_energy_job()
      result = do(gmsList,output=1,punch=1,
                  run_prefix=run_prefix)
      self.read_output_list(result[0])
      self.read_punch_list(result[1])

   def get_optimized_energy(self,run_prefix=None,zmat_extend=None):
      gmsList = self.get_optimize_job(zmat_extend=zmat_extend)
      result = do(gmsList,output=1,punch=1,
                  run_prefix=run_prefix)
      self.read_output_list(result[0])
      self.read_punch_list(result[1])

   def get_optimized_charges(self,run_prefix=None,skip=None):
      gmsList = self.get_optimize_charge_job()
      result = do(gmsList,output=1,punch=1,
                  run_prefix=run_prefix,skip=skip)
      self.read_output_list(result[0])
      self.read_punch_list(result[1])

   def get_prop_charges(self,run_prefix=None):
      gmsList = self.get_prop_charge_job()
      result = do(gmsList,output=1,punch=1,
                  run_prefix=run_prefix)
      self.read_output_list(result[0])
      self.read_punch_list(result[1])


if os.environ.has_key('GAMESS'):
   base = os.environ['GAMESS']
   rungms_path = base + '/rungms'
else:
   base = ''
   rungms_path = ''






