#A* -------------------------------------------------------------------
#B* This file contains source code for the PyMOL computer program
#C* copyright 1998-2000 by Warren Lyford Delano of DeLano Scientific. 
#D* -------------------------------------------------------------------
#E* It is unlawful to modify or remove this copyright notice.
#F* -------------------------------------------------------------------
#G* Please see the accompanying LICENSE file for further information. 
#H* -------------------------------------------------------------------
#I* Additional authors of this source file include:
#-* 
#-* 
#-*
#Z* -------------------------------------------------------------------

# 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 + '/bin/rungms'
else:
    base = ''
    rungms_path = ''






