#  Copyright (C) 2021 The authors of Py-ChemShell
#
#  This file is part of Py-ChemShell.
#
#  Py-ChemShell is free software: you can redistribute it and/or modify
#  it under the terms of the GNU Lesser General Public License as
#  published by the Free Software Foundation, either version 3 of the
#  License, or (at your option) any later version.
#
#  Py-ChemShell 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 Lesser General Public License for more details.
#
#  You should have received a copy of the GNU Lesser General Public
#  License along with Py-ChemShell.  If not, see
#  <http://www.gnu.org/licenses/>.

__author__ = ('Tom Keal <thomas.keal@stfc.ac.uk>')

from ...io   import file
from ...data import atomicdata, unitconvert
from ..qm    import _QM

from os import getpid
from shutil import copyfile

class Molpro(_QM):
    '''Interface to Molpro'''

    _kwargs = {
                'bq_grads'   :'_molpro.bqgrads',
                'jobname'    :'_molpro',
                'runargs'    :'_molpro.inp',
                'input'      :('_molpro.inp','_molpro.lat'),
                'output'     :'_molpro.out',
                'memory'     :256,
                'scf_threshold':6,
              }

    _internals = {
                   '_syscmd': 'molpro'
                 }

    # must be a procedure rather than a direct definition, otherwise frag.natoms is unknown
    def setOutKeys(self):
        '''Define {task:keyword} for searching in output file'''

        if self.frag.bqs:
            nbq = len(self.frag.bqs.coords)
        else:
            nbq = 0 
        ncent = self.frag.natoms + nbq

        # each keyword should contain an array of integers specifying:
        #     [ row_0, # starting row
        #       row_n, # ending row
        #       row_i, # row increment
        #       col_0, # starting col
        #       col_n, # ending col
        #       col_i, # col increment
        #       skip , # skip the previous (skip-1) match,
        #       coeff, # unit conversion coefficient
        #     ]
        # referencing to the line where the key is matched
        self._outkeys = {
                          'energy'   : {"ChemShell energy":[1,1,1,2,2,1,1,1.0]},
                          'gradients': {"^ Atom          dE/dx               dE/dy               dE/dz$":[2,2+self.frag.natoms-1,1,1,3,1,1,1.0]},
                          'bq_grads' : {"gradient output":[2,2+nbq-1,1,0,2,1,1,1.0]},
                        }


    def getInpBuff(self, input=input):
        '''String buffer containing the input information.'''

        # Generate input file as list of strings
        inp = []

        natoms = self.frag.natoms
        if self.frag.bqs:
            nbq = len(self.frag.bqs.coords)
        else:
            nbq = 0 
        ncent = natoms + nbq

        # Main input (.inp) file
        if input == self.input[0]:

            # Title
            inp.append("***,%s" % self.jobname)

            # Memory specification
            inp.append("memory,%s,m" % str(self.memory))

            # Restart file specification
            wfuName = self.jobname + str(getpid()) + ".wfu"
            if self.restart:
                inp.append("file,2,%s" % wfuName)
            else:
                inp.append("file,2,%s,new" % wfuName)

            # Integral-direct calculations
            if self.direct:
                inp.append("gdirect")

            # Threshold for SCF convergence
            inp.append("gthresh,energy=1.d-%s" % self.scf_threshold)

            # Spin multiplicity
            if self.mult > 1:
                inp.append("set,spin=%s" % str(self.mult-1))

            # Geometry
            inp.append("orient,noorient")
            inp.append("symmetry,nosym")
            inp.append("geomtyp=xyz")
            inp.append("bohr")
            inp.append("geometry={")
            inp.append("%s" % natoms)
            inp.append("%s" % self.frag.frag_name)
            for i in range(natoms):
                coords = list(self.frag.coords[i])
                z = self.frag.znums[i]
                symbol = str(atomicdata.symbols[z])
                inp.append("%s,%s,%s,%s" % (symbol, coords[0], coords[1], coords[2]))
            inp.append("}")

            # Charge
            if self.charge != 0:
                inp.append("set,charge=%s" % self.charge)

            # Point charge file
            if nbq > 0:
                inp.append("lattice,%s,%s" % (self.input[1], self.bq_grads))

            # Basis set
            if not self.harmonic:
                inp.append("cartesian")
            inp.append("basis=%s" % self.basis)

            # Wave function information
            if self.method == 'dft':
                if self.scftype == 'uhf' or self.scftype == 'uks':
                    inp.append("uks,%s" % self.functional)
                else:
                    inp.append("ks,%s" % self.functional)
            else:
                if self.scftype == 'uhf':
                    inp.append("uhf")
                else:
                    inp.append("rhf")

            # Print out energy for easy parsing
            inp.append("text,ChemShell energy")
            inp.append("show[5,5f20.12], energy")

            # Requested properties
            if self._gradients:
                inp.append("force")


        # Lattice of point charges for QM/MM calculations in a separate input file
        elif input == self.input[1] and nbq > 0:

            inp.append("Point Charges")
            inp.append("%s" % nbq)
            for i in range(nbq):
                # Lattice file requires coordinates in Angstrom
                coords = [x * unitconvert.Bohr2Angstrom for x in self.frag.bqs.coords[i]]
                charge = self.frag.bqs.charges[i]
                inp.append("%s %s %s %s 1" % (coords[0], coords[1], coords[2], charge))


        # Generate input buffer from input list
        buff = "\n".join(inp)
        return buff


    def parseMainOutput(self):
        '''Parse main output file self.output and bq gradient output file.'''

        from ...io import file

        self.setOutKeys()

        self._result.energy    = file.getArrayFromFile(self.output, 'out', self._result.energy.shape   , regex=self._outkeys['energy'])
        if self._gradients:
            self._result.gradients = file.getArrayFromFile(self.output, 'out', self._result.gradients.shape, regex=self._outkeys['gradients'])
        if self.frag.nbqs and self._gradients:
            self.frag.bqs.gradients = file.getArrayFromFile(self.bq_grads, 'out', self.frag.bqs.gradients.shape, regex=self._outkeys['bq_grads'])


