#  Copyright (C) 2017 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
from ..qm    import _QM

from shutil import copyfile

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

    _kwargs = { 
                'bq_grads'   :'_lsdalton.bq_grads',
                'guess'      :'ATOMS',
                'jobname'    :'_lsdalton',
                'runargs'    :'_lsdalton',
                'input'      :'_lsdalton.dal',
                'output'     :'_lsdalton.out',
                'auxbasis'   :'',
                'harmonic'   : True,
                'convdyn'    :'standard',
                'scalapack'  : False,
              }

    _internals = {
                   '_syscmd': 'lsdalton'
                 }

    # 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
        # YL 07/03/2021: initially the regex clause was \(au\) but newer Python will throw out
        #                "DeprecationWarning: invalid escape sequence"; recently the antislash
        #                escape was wrongly removed by me now the correct regex is used
        #                there might be more similar cases in the package but we'll fix
        #                when we see them again
        self._outkeys = {
                          'energy'   : {"Final HF energy:":[0,0,1,3,3,1,1,1.0],
                                        "Final DFT energy:":[0,0,1,3,3,1,1,1.0]},
                          'gradients': {'Molecular gradient \\(au\\)':[3,3+self.frag.natoms-1,1,1,3,1,1,1.0]},
                          'bq_grads' : {"Molecular gradient \\(au\\)":[3+self.frag.natoms,3+ncent-1,1,1,3,1,1,1.0]},
                        }


    def getInpBuff(self):
        '''String buff containing the input information

        For convenience we use the single-file input format,
        with the .mol data listed first followed by .dal'''

        # Check for options not supported by LSDalton
        # UHF not supported (See https://daltonprogram.org/features subsection SCF optimization)
        if self.scftype.lower() not in ['rhf', 'rks'] or self.mult > 1:
            print('\n ERROR: LSDalton only implements restricted HF and DFT for SCF optimisation')
            exit(999)

        # 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

        # The charge specified in the input file equals
        # the user-defined charge plus point charges, if any. 
        # (so the number of electrons stays constant regardless
        #  of changes to the MM environment)
        inpcharge = self.charge
        for i in range(nbq):
            inpcharge += self.frag.bqs.charges[i]

        # ============= .mol input ==============
        # Basis set
        if nbq > 0:
            #  use individual basis lines rather than $basis
            inp.append("ATOMBASIS")
            use_atombasis = True
        else:
            inp.append("BASIS")
            if not self.auxbasis:
                inp.append("%s" % self.basis)
            else:
                inp.append("%s Aux=%s" % (self.basis, self.auxbasis))
            use_atombasis = False
        # Title
        inp.append("ChemShell LSDALTON interface")
        inp.append("")
        # Geometry
        # Not grouping atoms so atomtypes = no of atoms
        if self.harmonic:
            harmout = "Spherical"
        else:
            harmout = "Cartesian"
        inp.append("Atomtypes=%s Charge=%s %s" % (ncent, inpcharge, harmout))
        for i in range(natoms):
            coords = list(self.frag.coords[i])
            z = self.frag.znums[i]
            symbol = str(atomicdata.symbols[z])
            if use_atombasis:
                atombasis = "Basis=%s" % self.basis
                if not self.auxbasis:
                    auxbasisout = ""
                else:
                    auxbasisout = "Aux=%s" % self.auxbasis
            else:
                atombasis = ""
                auxbasisout = ""
            inp.append("Charge=%s Atoms=1 %s %s" % (z, atombasis, auxbasisout))
            inp.append("%s %s %s %s" % (symbol, coords[0], coords[1], coords[2]))
        for i in range(nbq):
            coords = list(self.frag.bqs.coords[i])
            q = self.frag.bqs.charges[i]
            inp.append("Charge=%s Atoms=1 Basis=pointcharge %s" % (q, auxbasisout))
            inp.append("bq %s %s %s" % (coords[0], coords[1], coords[2]))
        inp.append("")
        #
        # ============= .dal input ==============
        # General settings
        inp.append("**GENERAL")
        if self.scalapack:
            inp.append(".SCALAPACK")
        # Calculation of integrals
        inp.append("**INTEGRAL")
        if self.frag.bqs:
            inp.append(".NOBQBQ")
        # Wave function information
        inp.append("**WAVE FUNCTIONS")
        if self.method == 'dft':
            inp.append(".DFT")
            inp.append("%s" % self.functional)
        else:
            inp.append(".HF")
        inp.append("*DENSOPT")
        inp.append(".CONVDYN")
        inp.append("%s" % self.convdyn.upper())
        inp.append(".MAXIT")
        inp.append("%d"%self.maxiter)
        if self.guess:
            inp.append(".START")
            inp.append(f"{self.guess}")
        if self.restart:
            inp.append(".RESTART")
            # TD 05/06/2024: Used in Tcl-ChemShell, this purify flag is required for geometry optimisations
            #                Without it they will crash if too big a geometry step is taken
            inp.append(".PURIFYRESTARTDENSITY")
        # Requested properties
        if self._gradients:
            inp.append("**RESPONS")
            inp.append("*MOLGRA")
        inp.append("*END OF INPUT")
        inp.append("")

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


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

        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._gradients and 'bq_grads' in self._outkeys and self.frag.nbqs:
            self.frag.bqs.gradients = file.getArrayFromFile(self.output, 'out', self.frag.bqs.gradients.shape, regex=self._outkeys['bq_grads'])


    def runLib(self, options=None, **kwargs):
        '''Run LSDalton as a shared library (called by run.runLib())'''

        from os         import path
        from ctypes     import c_int, byref, CDLL
        from chemsh     import CHEMSH_LIB_DIR
        from ...dl_py2f import py2f

        print("\n >>> Directly-linked LSDalton library will be used...")       

        # Replicate behaviour of lsdalton wrapper script
        copyfile(self.input, 'LSDALTON.INP')
        copyfile(self.input, 'MOLECULE.INP')
        
        # load LSDalton library: use cross-platform function os.path.join()
        soname = path.join(CHEMSH_LIB_DIR, 'liblsdalton.so')

        liblsdalton = CDLL(soname)

        liblsdalton.interface_lsdalton(byref(py2f(self.frag)))

        # Replicate behavour of lsdalton wrapper script
        copyfile('LSDALTON.OUT', self.output)
        
