#!/usr/bin/env python
"""
This script generates input files from XZY file provided as input based on the
available template file.

Usage: doinp.py [options] xyz file(s)

Options:
  -h, --help       show this help
  -s, --gaussian   prepare gaussian inputs (expects gaussian.tmpl)
  -g, --gamess     prepare gamess inputs (expects gamess.tmpl)
  -m, --molcas     prepare molcas inputs (expects molcas.tmpl)
      --solvents   prepare input files for a following set of solvents:
                   n-hexane (1.8819)
                   toluene (2.3741)
                   chlorobenzene (5.6968)
                   acetonitrile (35.688)
                   water (78.3553)
"""

#     Copyright (C) 2011, Robert W. Gora (robert.gora@pwr.wroc.pl)
#  
#     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.,
#     675 Mass Ave, Cambridge, MA 02139, USA.

__author__  = "Robert Gora (robert.gora@pwr.wroc.pl)"
__version__ = filter(str.isdigit, "$Revision: 1 $")

# Import necessary modules
import os, sys, getopt, re

from string import Template
#from numpy import *

# Regular expressions
reflags = re.DOTALL

#----------------------------------------------------------------------------
# Usage
#----------------------------------------------------------------------------
def Usage():
    """Print usage information and exit."""
    print __doc__

    sys.exit()

#----------------------------------------------------------------------------
# Main
#----------------------------------------------------------------------------
def Main(argv):
    '''Parse commandline and loop throught the logs'''

    data = {}

    # Set up defaults
    gamess = 0
    gaussian = 1
    molcas = 0
    solvents = []

    # Parse commandline
    try:
        opts, args = getopt.getopt(argv, "hsgm",
                                        ["help",
                                         "gaussian",
                                         "gamess",
                                         "molcas",
                                         "solvents"
                                         ])
    except getopt.GetoptError, error:
        print(error)
        Usage()
    if not argv:
        Usage()

    for opt, arg in opts:
        if opt in ("-h", "--help"):
            Usage()
        elif opt in ("-g", "--gamess"):
            gamess=1
        elif opt in ("-s", "--gaussian"):
            gaussian=1
        elif opt in ("-m", "--molcas"):
            molcas=1
        elif opt in ("--solvents"):
            #solvents=['n-hexane', 'toluene', 'chlorobenzene', 'acetonitrile', 'water', 'dmso', 'n,n-diMethylformamide', 'methanol', 'acetone', 'n-octanol', 'chloroform', 'carbontetrachloride', 'tetrahydrofuran', 'tetrahydrofuran']
            solvents=['n-hexane', 'toluene', 'chlorobenzene', 'acetonitrile', 'water', 'dmso', 'n,n-diMethylformamide', 'methanol', 'acetone', 'n-octanol', 'chloroform', 'carbontetrachloride', 'tetrahydrofuran', 'tetrahydrofuran']

    # Parse each data file (with xyz coords)
    data_files = args

    # Prepare inputs
    for f in data_files:
        if gaussian:
            GAUSSIAN_INPUTS(f,solvents)
        if gamess:
            GAMESS_INPUTS(f)
        if molcas:
            MOLCAS_INPUTS(f)


#----------------------------------------------------------------------------
# Common input template routines
#----------------------------------------------------------------------------
class INPUT_TEMPLATE(Template):
    delimiter = '@'

class INPUTS:
    """Common input routines"""
    def __init__(self, data):
        self.data = data

        self.ReadTemplate()

    def ReadTemplate(self):
        """Read or punch standard template"""
        try:
            self.tmpl = open(self.pkg+'.tmpl','r').read()
            self.tmpl = INPUT_TEMPLATE(self.tmpl)
            self.WriteInputs()
        except IOError:
            print "There's no " + self.pkg + " template. I'm punching one - please check"
            open(self.pkg+'.tmpl','w').write(self.tmpl)
            sys.exit()
        except AttributeError:
            pass

    def WriteInputs(self):
        pass

#----------------------------------------------------------------------------
# Gamess (US) routines
#----------------------------------------------------------------------------

class GAMESS_INPUTS(INPUTS):
    """Gamess US input routines"""

    def __init__(self, data):
        # template name
        self.pkg = "gamess"
        # template content
        self.tmpl="""\
 $system mwords=10 memddi=0 parall=.t. $end
 $contrl scftyp=rhf runtyp=energy icharg=0 mult=1 units=angs
         maxit=100 nprint=8 exetyp=run mplevl=0 ispher=-1
         icut=20 itol=30 aimpac=.f. cctyp=none $end
!$eds    mch(1)=0,0 mmul(1)=1,1 mnr(1)=1 ffeds=.t. $end
!$mp2    mp2prp=.t. code=ddi cutoff=1d-20 $end
!$ccinp  iconv=14 maxcc=100 $end
!$trans  cuttrf=1d-15 $end
 $scf    dirscf=.t. fdiff=.f. diis=.t. soscf=.f.
         conv=1d-11 swdiis=0.0001 npreo(1)=1,-1,2,1 $end
 $basis  gbasis=sto ngauss=3 $end
@data
"""
        INPUTS.__init__(self, data)

    def WriteInputs(self):

        # initialize periodic table
        p=Periodic(0)
        
        # read xyz file
        try:
            xyz=open(self.data,'r').readlines()
            xyz=xyz[2:int(xyz[0])+2]
        except ValueError:
            print "Problem with *.xyz file?"
            sys.exit(1)
        
        for i in range(len(xyz)):
            xyz[i]=xyz[i].split()
            xyz[i].insert(1, str( Atomn(xyz[i][0], p) ))
            xyz[i]='%-5s %5s %15s %15s %15s\n' % tuple([xyz[i][j] for j in range(5)])
        
        xyz.insert(0, ' $data\n%s\nc1 0\n' % self.data)
        xyz.append(' $end')
        xyz=''.join(xyz)
        
        # write input files
        filename = self.data.replace('.xyz','')
        filename = filename.replace(' ','_')
        finput = self.tmpl.substitute(data=xyz)

        open(filename+'.inp','w').write(finput)

#----------------------------------------------------------------------------
# Molcas routines
#----------------------------------------------------------------------------

class MOLCAS_INPUTS(INPUTS):
    """Molcas input routines"""

    def __init__(self, data):
        # template name
        self.pkg = "molcas"
        # template content
        self.tmpl="""\
 &GATEWAY
Coord
@data
Basis=STO-3G
Title=?
Group=NoSym

 &SEWARD

 &SCF
thre     = 1.0e-15, 1.0e-10, 1.0e-08, 1.0e-07

 &MBPT2

 &RASSCF
outorb   = canonical

 &MOTRA
frozen = 0

 &CCSDT
CCT
iterations
200
accuracy = 1.0e-15

 &RASSCF
nactel   = ? 0 0
inactive = ?
ras2     = ?
thrs     = 1.0e-15, 1.0e-06, 1.0e-06

 &CASPT2

"""
        INPUTS.__init__(self, data)

    def MakeCoords(self):
        '''read xyz file'''
        try:
            xyz=open(self.data,'r').readlines()
            xyz=xyz[0:int(xyz[0])+2]
            xyz=''.join(xyz).rstrip()
        except ValueError:
            print "Problem with *.xyz file?"
            sys.exit(1)
        except IndexError:
            print "Empty *.xyz file - assuming old molcas syntax"
            xyz=''

        return xyz

    def WriteInputs(self):

        # read xyz file
        xyz=self.MakeCoords()

        # write input files
        filename = self.data.replace('.xyz','')
        filename = filename.replace(' ','_')

        finput = self.tmpl.substitute(data=xyz)
        open(filename+'.inp','w').write(finput)

#----------------------------------------------------------------------------
# Gaussian routines
#----------------------------------------------------------------------------

class GAUSSIAN_INPUTS(INPUTS):
    """Gaussian 09 input routines"""

    def __init__(self, data, solvents):
        # template name
        self.pkg = "gaussian"
        # solvents
        self.solvents = solvents
        # template content
        self.tmpl="""\
# m06l/6-31g(d,p) nosymm opt=vtight freq int=ultrafine scrf(cpcm, solvent=@solvent)

gaussian input

0 1
@data
"""
        INPUTS.__init__(self, data)

    def WriteInputs(self):

        # read xyz file
        try:
            xyz=open(self.data,'r').readlines()
            xyz=xyz[2:int(xyz[0])+2]
            xyz=''.join(xyz)
        except ValueError:
            print "Problem with *.xyz file?"
            sys.exit(1)

        if self.solvents:
            # write input files
            for s in self.solvents:
                filename = self.data.replace('.xyz','')+'_%s' % s
                filename = filename.replace(' ','_')
                finput = self.tmpl.substitute(data=xyz, solvent=s)
                open(filename+'.com','w').write(finput)
        else:
            filename = self.data.replace('.xyz','')
            filename = filename.replace(' ','_')
            finput = self.tmpl.substitute(data=xyz)
            open(filename+'.com','w').write(finput)

#----------------------------------------------------------------------------
# Utilities
#----------------------------------------------------------------------------

def Run(program, args):
    '''Run shell command'''
    os.system(program + ' ' + args)

def SkipLines(open_file,n):
    '''Read n lines from file f.'''

    for i in range(n):
        line = open_file.readline()
        if line == '' :
            break

    return line

def FindLine(open_file,pattern):
    '''Read lines until pattern matches.'''

    while 1:
        line = open_file.readline()
        if line.find(pattern) !=-1 :
            break
        if line == '' :
            line = -1
            break

    return line

def ReFindLine(open_file,pattern):
    '''Read lines until pattern matches.'''

    while 1:
        line = open_file.readline()
        if re.search(pattern,line) !=None : break
        if line == '' :
            linie = -1
            break

    return line

def Periodic(mendeleiev):
    '''Returns the mendeleiev table as a python list of tuples. Each cell
    contains either None or a tuple (symbol, atomic number), or a list of pairs
    for the cells * and **. Requires: "import re". Source: Gribouillis at
    www.daniweb.com - 2008 '''

    # L is a consecutive list of tuples ('Symbol', atomic number)
    L = [ (e,i+1) for (i,e) in enumerate( re.compile ("[A-Z][a-z]*").findall('''
    HHeLiBeBCNOFNeNaMgAlSiPSClArKCaScTiVCrMnFeCoNiCuZnGaGeAsSeBrKr
    RbSrYZrNbMoTcRuRhPdAgCdInSnSbTeIXeCsBaLaCePrNdPmSmEuGdTbDyHoEr
    TmYbLuHfTaWReOsIrPtAuHgTlPbBiPoAtRnFrRaAcThPaUNpPuAmCmBkCfEsFm
    MdNoLrRfDbSgBhHsMtDsRgUubUutUuqUupUuhUusUuo'''))]

    # The following fills the void with nones and returns the list of lists
    mendeleiev = 0

    if mendeleiev:
        for i,j in ( (88,103), (56,71) ):
            L[i] = L[i:j]
            L[i+1:] = L[j:]
        for i,j in ( (12,10 ), (4,10), (1,16) ):
            L[i:i]=[None]*j 

        return [ L[18*i:18*(i+1)] for i in range (7) ]

    # Return a plain list of tuples
    else:
        return L

def Atomn(s,ptable):
    '''Returns the atomic number based on atomic symbol string
    ptable is a list of consecutive (symbol, atomic number) tuples.'''

    for n,a in enumerate(ptable):
        if a[0].lower().find(s.strip().lower()) !=-1 :
            return float(n+1)

#----------------------------------------------------------------------------
# Main routine
#----------------------------------------------------------------------------
if __name__ == "__main__":
    Main(sys.argv[1:])