#  Copyright (C) 2020 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 ..qm    import _QM
from ...data import atomicdata, unitconvert

from shutil import copyfile

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

    _kwargs = {
                'jobname'    :"_mndo",
                'runargs'    :"<",
                'input'      :"_mndo.in",
                'output'     :"_mndo.out",
                'method'     : "mndo",
                'accuracy'   : "medium",
                'name'       : __name__,
              }

    _internals = {
                   '_syscmd': 'mndo'
                 }

    # 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

        # conversion factors from Tcl-ChemShell MNDO parser in interface_gamess/interface.c
        econv = 1.0 / 627.511525
        gconv = 1.0 / 1185.8253
        self._outkeys = {
            "energy"   : {"ENERGY, CARTESIAN AND INTERNAL GRADIENT NORM" : [1,1,0,0,0,1,1,econv]},
            "gradients": {"CARTESIAN GRADIENT$":[1,1+self.frag.natoms-1,1,2,4,1,1,gconv]},
            "bq_grads" : {"CARTESIAN GRADIENT OF MM ATOMS":[1,1+nbq-1,1,2,4,1,1,gconv]},
        }


    def getInpBuff(self):
        '''String buff 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

        # Keyword input lines

        # QM charge
        inpcharge = self.charge
        khargestr = "kharge=%s" % inpcharge 
            
        # Methods
        iop = {
            "mndo/d" : -10,
            "pm3" : -7,
            "am1" : -2,
            "scc" : 5,
            "mndo" : 0,
            "mindo/3" : 1,
            "cndo/2" : 2,
            "om1" : -5,
            "om2" : -6,
            "om3" : -8
        }
        iopstr = "iop=%s" % iop[self.method]

        # SCF type / multiplicity
        # scftype   mult   imult  iuhf
        #-------------------------------
        # rhf       1      undef  undef
        # rhf       N      N      -1
        # rohf      1      1      -1
        # rohf      N      N      -1
        # uhf       1      1       1
        # uhf       N      N       1
        
        scfstr = ""
        if self.scftype == 'rhf':
            if self.mult != 1:
                scfstr = "imult=%s iuhf=-1" % self.mult
        elif self.scftype == 'rohf':
            scfstr = "imult=%s iuhf=-1" % self.mult
        elif self.scftype == 'uhf':
            scfstr = "imult=%s iuhf=1" % self.mult

        # Electronic restarts
        trialstr = ""
        if self.restart:
            if self.method == "scc":
                trialstr = "ktrial=41"
            else:
                trialstr = "ktrial=11"

        # SCF accuracy
        accuracystr = ""
        if self.accuracy == "high":
            accuracystr = "iscf=8 iplscf=8"
        elif self.accuracy == "veryhigh":
            accuracystr = "iscf=10 iplscf=10"

        # Electrostatic embedding options
        # Currently only the standard "shift" method is implemented.
        # (as opposed to L1 etc. which are supported by Tcl-ChemShell
        # but I suspect rarely used)
        if nbq:
            mminpstr = "mminp=2"
            mmcoupstr = "mmcoup=2"
            mmlinkstr = "mmlink=2"
            # MNDO inputs require a list of link atoms when using electrostatic embedding
            # (NB apparently not required when using mechanical embedding, at least
            # Tcl-ChemShell does not provide them in this case)
            nlink = len(self._linkatoms)
            nlinkstr = "nlink=%s" % nlink
            numatmstr = "numatm=%s" % nbq
        # Calculation mode
        if self._gradients:
            jopstr = "jop=-2"
        else:
            jopstr = "jop=-1"

        # Number of SCF iterations
        iterstr = "kitscf={}".format(self.maxiter)
        
        # Build input keyword lines
        inp.append("%s %s %s %s +" % (khargestr, iopstr, scfstr, trialstr))
        inp.append("%s %s +" % (accuracystr, iterstr))
        # The following options are mostly fixed to follow the input format we are generating
        # TODO implement binary output option (cf Tcl-ChemShell)
        inp.append("igeom=1 iform=1 nsav15=4 ipubo=1 +")
        #  Electrostatic embedding
        if nbq:
            inp.append("%s %s %s %s %s +" % (mminpstr, mmcoupstr, mmlinkstr, nlinkstr, numatmstr))
        inp.append("%s" % jopstr)

        inp.append("ChemShell MNDO interface:")
        inp.append("")

        # Geometry
        for i in range(natoms):
            # convert to angstrom
            coords = [ x * unitconvert.Bohr2Angstrom for x in self.frag.coords[i]]
            z = self.frag.znums[i]
            inp.append("%d %14.8f %d %14.8f %d %14.8f %d" % (z, coords[0], 0, coords[1], 0, coords[2], 0))

        # Termination of coordinates
        inp.append("0 0.0 0 0.0 0 0.0 0 ")

        # Electrostatic embedding input
        if nbq:
            # List of link atom indices
            # (note offset because MNDO counts from 1)
            if nlink:
                linkatmstr=""
                # Break lines at 80 characters
                for i in range(nlink):
                    linkatmstr += "%5d" % (self._linkatoms[i] + 1)
                    if not (i+1)%16 and (i+1) != nlink:
                        linkatmstr += "\n"
                inp.append(linkatmstr)
            # External point charge coordinates
            for i in range(nbq):
                # convert to angstrom
                coords = [ x * unitconvert.Bohr2Angstrom for x in self.frag.bqs.coords[i]]
                q = self.frag.bqs.charges[i]
                inp.append("%14.8f %14.8f %14.8f %9.6f" % (coords[0], coords[1], coords[2], q))

        # End of input
        inp.append("")

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


    def parseAdditionalOutput(self):
        '''Parse MNDO output file fort.15'''

        from ...io import file

        output = "fort.15"

        self.setOutKeys()

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

