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

from shutil import copyfile

class DFTBplus(_QM):
    '''Interface to DFTB+'''

    _kwargs = {
                'dispersion_param'     : {},
                'jobname'              :"_dftbplus",
                'runargs'              :"",
                'input'                :"dftb_in.hsd",
                'hubbard_derivs'       : {},
                'hydrogen_correction'  : '',
                'hydrogen_damping'     : '',
                'output'               :"dftb_in.out",
                'max_angular_momentum' : {},
                'method'               : 'mio',
                'name'                 : __name__,
                'scc'                  : True,
                'smear'                : True,
                'smear_method'         :'Fermi',
                'smear_mp_order'       : 2,   
                'smear_temperature'    : '0.0',
                'spin_constants'       : {},
                'skf_path'             : "",
                'third_order'          : None,
                'threshold'            : 1e-5,
              }

    _internals = {
                   '_syscmd': 'dftb+'
                 }

    # 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'   : {'Total energy:'             :[0,0,1,2,2,1,1,1.0]},
                          # YL 08/10/2020: updated according to DFTB+ v20.1
                          # NB: DFTB+ outputs "forces" rather than "gradients"
                          'gradients': {'Total Forces'              :[1,1+self.frag.natoms-1,1,1,3,1,1,-1.0]},
                          'bq_grads' : {'Forces on external charges':[1,1+nbq-1,1,0,2,1,1,-1.0]},
        }


    def getInpBuff(self):
        '''String buff containing the input information.'''
        
        # Setup default parameters based on the high-level method option
        # Most parameters from dftb.org/parameters
        # Default D3 dispersion correction parameters taken from DFTB+ manual, Appendix G
        default_angular_momentum = {}
        default_hubbard_derivs = {}
        default_hydrogen_damping = 0.0
        default_spin_constants = {} 
        # DFTB+ D3 "old default" parameters
        default_dispersion_param = {'a1':0.5719, 'a2':3.6017, 's6':1.0, 's8':0.5883}
        default_third_order = False

        # Values taken from DFTB+ manual, Appendix E
        spin_constants_LDA = {
                 #    Wss
                'H' :'-0.064',
                 #    Wss     Wsp     Wps     Wpp 
                'C' :'-0.028  -0.024  -0.024  -0.022',
                 #    Wss     Wsp     Wps     Wpp 
                'N' :'-0.030  -0.026  -0.026  -0.025',
                 #    Wss     Wsp     Wps     Wpp 
                'O' :'-0.032  -0.028  -0.028  -0.027',
                 #    Wss     Wsp     Wsd     Wps     Wpp     Wpd     Wds     Wdp     Wdd
                'Si':'-0.018  -0.013   0.000  -0.013  -0.012  0.000   0.000   0.000   -0.019',
                 #    Wss     Wsp     Wsd     Wps     Wpp     Wpd     Wds     Wdp     Wdd
                'S' :'-0.019, -0.016,  0.000, -0.016, -0.014,  0.000, 0.000,   0.000, -0.010',
                 #    Wss     Wsp     Wsd     Wps     Wpp     Wpd     Wds     Wdp     Wdd
                'Fe':'-0.013, -0.009, -0.003, -0.009, -0.011, -0.001, -0.003, -0.001, -0.015',
                 #    Wss     Wsp     Wsd     Wps     Wpp     Wpd     Wds     Wdp     Wdd
                'Ni':'-0.009, -0.009, -0.003, -0.009, -0.010, -0.001, -0.003, -0.001, -0.017'
                }
        # Values taken from DFTB+ manual, Appendix E
        spin_constants_PBE = {
                 #    Wss
                'H' :'-0.072',
                 #    Wss     Wsp     Wps     Wpp 
                'C' :'-0.031  -0.025  -0.025  -0.023',
                 #    Wss     Wsp     Wps     Wpp 
                'N' :'-0.033  -0.027  -0.027  -0.026',
                 #    Wss     Wsp     Wps     Wpp 
                'O' :'-0.035  -0.030  -0.030  -0.028',
                 #    Wss     Wsp     Wsd     Wps     Wpp     Wpd     Wds     Wdp     Wdd
                'Si':'-0.020  -0.015   0.000  -0.015  -0.014   0.000   0.002   0.002  -0.032',
                 #    Wss     Wsp     Wsd     Wps     Wpp     Wpd     Wds     Wdp     Wdd
                'S' :'-0.021  -0.017   0.000  -0.017  -0.016   0.000   0.000   0.000  -0.080',
                 #    Wss     Wsp     Wsd     Wps     Wpp     Wpd     Wds     Wdp     Wdd
                'Fe':'-0.016  -0.012  -0.003  -0.012  -0.029  -0.001  -0.003  -0.001  -0.015',
                 #    Wss     Wsp     Wsd     Wps     Wpp     Wpd     Wds     Wdp     Wdd
                'Ni':'-0.016  -0.012  -0.003  -0.012  -0.022  -0.001  -0.003  -0.001  -0.018'
                }

        # Parameter set "3ob-3-1"
        if self.method.lower().strip() == '3ob':
            default_angular_momentum = {'Br':'d', 'C':'p', 'Ca':'p', 'Cl':'d', 'F':'p', 'H':'s', 'I':'d', 'K':'p', 'Mg':'p', 'N':'p', 'Na':'p', 'O':'p', 'P':'d', 'S':'d', 'Zn':'d'}
            if self.hydrogen_correction == '':
                self.hydrogen_correction = 'damping'
            default_hydrogen_damping = 4.0
            default_third_order = True
            default_hubbard_derivs = {'Br':-0.0573, 'C':-0.1492, 'Ca':-0.0340, 'Cl':-0.0697, 'F':-0.1623, 'H':-0.1857, 'I':-0.0433, 'K':-0.0339,
                                     'Mg':-0.02, 'N':-0.1535, 'Na':-0.0454, 'O':-0.1575, 'P':-0.14, 'S':-0.11, 'Zn':-0.03}
            default_dispersion_param = {'a1':0.746, 'a2':4.191, 's6':1.0, 's8':3.209}
            #J. Chem. Theory Comput. 2013, 9, 1, 338–354
            default_spin_constants = spin_constants_PBE

        # Parameter set "mio-1-1"
        elif self.method.lower().strip() == 'mio':
            default_angular_momentum = {'H':'s', 'C':'p', 'N':'p', 'O':'p', 'P':'d', 'S':'d'}
            default_third_order = False
            # Phys. Rev. B 58, 7260
            default_spin_constants = spin_constants_LDA

        # Parameter set "ob2-1-1" (has three versions)
        elif self.method.lower().strip() == 'ob2_base':
            default_angular_momentum = {'H':'s', 'C':'p', 'N':'p', 'O':'p'}
            default_third_order = False
            default_dispersion_param = {'a1':0.717, 'a2':2.565, 's6':1.0, 's8':0.011}
            # J. Chem. Theory Comput. 2018, 14, 1, 115–125
            default_spin_constants = spin_constants_PBE
        elif self.method.lower().strip() == 'ob2_shift':
            default_angular_momentum = {'H':'s', 'C':'p', 'N':'p', 'O':'p'}
            default_third_order = False
            default_dispersion_param = {'a1':0.816, 'a2':2.057, 's6':1.0, 's8':0.010}
            # J. Chem. Theory Comput. 2018, 14, 1, 115–125
            default_spin_constants = spin_constants_PBE
        elif self.method.lower().strip() == 'ob2_split':
            default_angular_momentum = {'H':'s', 'C':'p', 'N':'p', 'O':'p'}
            default_dispersion_param = {'a1':0.497, 'a2':3.622, 's6':1.0, 's8':0.010}
            default_third_order = False
            # J. Chem. Theory Comput. 2018, 14, 1, 115–125
            default_spin_constants = spin_constants_PBE

        # Parameter set "matsci-0-3"
        elif self.method.lower().strip() == 'matsci':
            default_angular_momentum = {'Al':'d', 'B':'p', 'C':'p', 'Cu':'d', 'H':'s', 'Na':'p', 'N':'p', 'O':'p', 'P':'d', 'Si':'d', 'Ti':'d'}
            default_third_order = False
            default_spin_constants = spin_constants_PBE

        # Parameter set "pbc-0-3"
        elif self.method.lower().strip() == 'pbc':
            # I don't think it is safe to auto-set pbc angular momentum as the  recommended Si value changes (Si:p for Silicon and Si:d for Silica)
            #    default_angular_momentum = {}
            default_third_order = False
            default_spin_constants = spin_constants_LDA

        # If third order hasn't been set explicitly set to default
        if self.third_order is None:
            self.third_order = default_third_order

        # 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]
            
        # Geometry (GenFormat)
        inp.append("Geometry = GenFormat {")
        # C = cluster (non-periodic) calculation
        if self.frag.cell.ndimensions == 0:
            inp.append("%s C" % natoms)
        # S = supercell (periodic) calculation
        elif self.frag.cell.ndimensions == 3:
            inp.append("%s S" % natoms)
        else:
            print("\n ChemShell/DFTB+ ERROR: Fragment dimensionality not supported \"{}\"".format(self.frag.cell.ndimensions))
            exit(112)
        # Build a list of elements in the system and a mapping from znums to the list
        element_list = []
        znum_list = []
        znum_map = {}
        for i in range(natoms):
            z = self.frag.znums[i]
            if z not in znum_map:
                znum_list.append(z)
                element_list.append(str(atomicdata.symbols[z]))
                znum_map[z] = len(element_list) - 1
        elements = " ".join(element_list)
        inp.append(elements)
        # Coordinates
        for i in range(natoms):
            # convert to angstrom
            coords = [ x * unitconvert.Bohr2Angstrom for x in self.frag.coords[i]]
            z = self.frag.znums[i]
            # Types count from 1 in the input
            type = znum_map[z] + 1
            inp.append("    %s %s %s %s %s" % (i+1, type, coords[0], coords[1], coords[2]))
        # Write lattice constants (if required)
        # C = cluster (non-periodic) calculation
        if self.frag.cell.ndimensions == 0:
            pass
        # S = supercell (periodic) calculation
        elif self.frag.cell.ndimensions == 3:
            inp.append("    %s %s %s" % (0.0, 0.0, 0.0))
            inp.append("    %s %s %s" % (self.frag.cell.vectors[0,0], self.frag.cell.vectors[0,1], self.frag.cell.vectors[0,2]))
            inp.append("    %s %s %s" % (self.frag.cell.vectors[1,0], self.frag.cell.vectors[1,1], self.frag.cell.vectors[1,2]))
            inp.append("    %s %s %s" % (self.frag.cell.vectors[2,0], self.frag.cell.vectors[2,1], self.frag.cell.vectors[2,2]))
        else:
            print("\n ChemShell/DFTB+ ERROR: Fragment dimensionality not supported \"{}\"".format(self.frag.cell.ndimensions))
            exit(112)
        inp.append("}")

        # Empty Driver = single point calculation
        inp.append("Driver = {}")

        # Hamiltonian
        inp.append("Hamiltonian = DFTB {")
        if self.scc:
            inp.append("  SCC = Yes")
            inp.append("  SCCTolerance = {}".format(self.threshold))
            inp.append("  MaxSCCIterations = %s" % self.maxiter)
        
        # Spin polarisation
        # Use defaults, but replace with user-supplied values if given
        self.spin_constants = default_spin_constants | self.spin_constants
        if self.scftype.lower() in [ 'uhf', 'uks' ]:
            inp.append("  SpinPolarisation = Colinear {")
            inp.append("    UnpairedElectrons = {}".format(self.mult - 1))
            inp.append("    RelaxTotalSpin = No")
            inp.append("  }")
            inp.append("  SpinConstants = {")
            inp.append("    ShellResolvedSpin = Yes")
            try:
                for i in range(len(element_list)):
                    element = element_list[i]
                    spin_constants = self.spin_constants[element]
                    inp.append('    {0} = {1}'.format(element, self.spin_constants[element]))
            except KeyError:
                print("\n ChemShell/DFTB+ ERROR: Missing spin constants for element \"{}\"".format(element))
                exit(112)
            inp.append("  }")
        elif self.scftype.lower() in [ 'rohf', 'roks' ]:
            pass

        # Use defaults, but replace with user-supplied values if given
        self.hubbard_derivs = default_hubbard_derivs | self.hubbard_derivs
        if self.third_order:
            inp.append("  ThirdOrderFull = Yes")
            inp.append("  HubbardDerivs {")
            try:
                for i in range(len(element_list)):
                    element = element_list[i]
                    hubbard_derivs = self.hubbard_derivs[element]
                    inp.append('    {0} = {1}'.format(element, hubbard_derivs))
            except KeyError:
                print("\n ChemShell/DFTB+ ERROR: Missing Hubbard U derivative value for element \"{}\"".format(element))
                exit(112)
            inp.append("  }")
        # Assuming here that the Slater-Koster filenames use the standard format
        inp.append("  SlaterKosterFiles = Type2FileNames {")
        if self.skf_path:
            inp.append('    Prefix = "%s"' % self.skf_path)
        inp.append('    Separator = "-"')
        inp.append('    Suffix = ".skf"')
        inp.append("  }")
        # Use defaults, but replace with user-supplied values if given
        self.max_angular_momentum = default_angular_momentum | self.max_angular_momentum
        inp.append("  MaxAngularMomentum {")
        try:
            for i in range(len(element_list)):
                element = element_list[i]
                max_angular_momentum = self.max_angular_momentum[element]
                inp.append('    {} = "{}"'.format(element, max_angular_momentum))
        except KeyError:
            print("\n ChemShell/DFTB+ ERROR: Missing angular momentum value for element \"{}\"".format(element))
            exit(112)

        inp.append("  }")

        # K-points
        if len(self.kpoints) == 3:
            inp.append("  KPointsAndWeights = SupercellFolding {")
            inp.append("    {0}  0   0  ".format(self.kpoints[0]))
            inp.append("     0  {0}  0  ".format(self.kpoints[1]))
            inp.append("     0   0  {0} ".format(self.kpoints[2]))
            # Check which k-pints are odd or even
            if (self.kpoints[0]%2 == 0):
                # l1 is even
                k_s1 = 0.5
            else:
                # l1 is odd
                k_s1 = 0.0
            if (self.kpoints[1]%2 == 0):
                # l1 is even
                k_s2 = 0.5
            else:
                # l1 is odd
                k_s2 = 0.0
            if (self.kpoints[2]%2 == 0):
                # l1 is even
                k_s3 = 0.5
            else:
                # l1 is odd
                k_s3 = 0.0
            inp.append("    {0} {1} {2} ".format(k_s1, k_s2, k_s3))
            inp.append("  }")

        # Should the density be restarted? (From charges.bin)
        if self.restart == False:
            inp.append("  ReadInitialCharges = No")
        elif self.restart == True:
            inp.append("  ReadInitialCharges = Yes")

        # Electronic smearing (called filling by DFTB+)
        if self.smear:
            smear_methods = [ "fermi", "methfesselpaxton" ]
            if self.smear_method.lower().strip() not in smear_methods:
                print("\n ChemShell/DFTB+ ERROR: SCC smearing method \"{}\" not valid. Available methods are:\n {}".format(self.smear_method.upper().strip(), smear_methods))
                exit(112)
            inp.append("  Filling = {} {{".format(self.smear_method.lower().strip()))
            inp.append("    Temperature [Kelvin] = {}".format(self.smear_temperature))
            if self.smear_method.lower().strip() == 'methfesselpaxton':
                if self.smear_mp_order < 1:
                    print("\n ChemShell/DFTB+ ERROR: Invalid Methfessel-Paxton order {} (must be greater than zero)".format(self.smear_mp_order))
                    exit(112)
                inp.append("    Order = {}".format(self.smear_mp_order))
            inp.append("  }")

        # Bqs
        if nbq > 0:
            inp.append("  ElectricField = {")
            inp.append("    PointCharges = {")
            inp.append("      CoordsAndCharges = {")
            for i in range(nbq):
                coords = list(self.frag.bqs.coords[i])
                q = self.frag.bqs.charges[i]
                inp.append("        %s %s %s %s" % (coords[0], coords[1], coords[2], q))
            inp.append("      }")
            inp.append("    }")
            inp.append("  }")
        
        # Dispersion corrections
        dispersion_methods = [ "", "lennardjones", "d3", "mbd"]
        # Use defaults, but replace with user-supplied values if given
        self.dispersion_param = default_dispersion_param | self.dispersion_param
        if self.dispersion.lower().strip() not in dispersion_methods:
                print("\n ChemShell/DFTB+ ERROR: Dispersion correction method \"{}\" not valid. Available methods are:\n {}".format(self.dispersion.lower().strip(), dispersion_methods))
                exit(112)
        # TD 16/2/24 Limited to Universal Force Field parameters for now
        # TODO allow manual parameter specification
        if self.dispersion.lower().strip() == "lennardjones":
            inp.append("  Dispersion = LennardJones {")
            inp.append("    Parameters = UFFParmeters {}")
            inp.append("  }")
        # TD 16/2/24 Dispersion model not yet implemented 
        elif self.dispersion.lower().strip() == "slaterkirkwood":
            pass
        elif self.dispersion.lower().strip() == "d3":
            inp.append("  Dispersion = SimpleDftD3 {")
            inp.append("    a1 = {}".format(self.dispersion_param['a1']))
            inp.append("    a2 = {}".format(self.dispersion_param['a2']))
            inp.append("    s6 = {}".format(self.dispersion_param['s6']))
            inp.append("    s8 = {}".format(self.dispersion_param['s8']))
            inp.append("  }")
        # TD 16/2/24 Dispersion model not yet implemented 
        elif self.dispersion.lower().strip() == "dftd4":
            pass
        elif self.dispersion.lower().strip() == "ts":
            inp.append("  Dispersion = TS {")
            inp.append("    RangeSeparation = 0.94")
            inp.append("  }")
        elif self.dispersion.lower().strip() == "mbd":
            inp.append("  Dispersion = MBD {")
            # Check if kpoints defined
            if len(self.kpoints) == 3:
                inp.append("  KGrid = {0} {1} {2}".format(self.kpoints[0], self.kpoints[1], self.kpoints[2]))
            # If not, default to Gamma-point sampling
            else:
                inp.append("  KGrid = 1 1 1")
            inp.append("  }")

        # Hydrogen corrections
        H_correction_methods = [ "", "damping"]

        if self.hydrogen_correction.lower().strip() not in H_correction_methods:
                print("\n ChemShell/DFTB+ ERROR: Hydrogen correction method \"{}\" not valid. Available methods are:\n {}".format(self.hydrogen_correction.lower().strip(), H_correction_methods))
                exit(112)
        if self.hydrogen_correction.lower().strip() == "damping":
            # If no user-supplied value given, use default
            if self.hydrogen_damping == '':
                self.hydrogen_damping = default_hydrogen_damping
            inp.append("  HCorrection = Damping {")
            inp.append("    Exponent = {}".format(self.hydrogen_damping))
            inp.append("  }")

        # end of Hamiltonian
        inp.append("}")

        inp.append("Options {}")

        inp.append("Analysis = {")
        if self._gradients:
            inp.append("  CalculateForces = Yes")
        else:
            inp.append("  CalculateForces = No")
        inp.append("}")

        # TD 19/2/24 Present ParserVersion 5 corresponds to DFTB+ v17.1
        # (Consult Appendix A of the DFTB+ manual for ParserVersion information)
        # Increasing this will break compatibility with old versions of DFTB+, but will update default parameters 
        # List of known version issues:
        # - CalculateForces has been renamed PrintForces, the old keyword stopped being supported in ParserVersion 14
        inp.append("ParserOptions {")
        inp.append("  ParserVersion = 5")
        inp.append("}")
        inp.append("")
        
        # Generate input buffer from input list
        buff = "\n".join(inp)
        return buff


    def parseAdditionalOutput(self):
        '''Parse detailed DFTB+ output file detailed.out'''

        from ...io import file

        output = "detailed.out"

        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'])



