#  Copyright (C) 2019 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__ = ['Andrew Logsdail <LogsdailA@cardiff.ac.uk>']

import chemsh
from ...data import unitconvert
from ..qm import _QM
from ...dl_py2f import utils
from ...parallel import parallel
from ...base import errors

class FHIaims(_QM):
    '''Interface to FHI-aims'''


    _kwargs = { 
               'input'            : ['control.in','geometry.in'],
               'initial_moment'   : {},
               'initial_charge'   : {},
               'ghost_species'    : None,
               'ghost_max_cutoff' : 5, #Setting this means we check by default for BQ-QM overlap, which is safest practice. Value is taken in Bohr.
               'ghost_min_cutoff' : 0.001, #This is used to determine if BQ-BQ interactions occur where the empty sites might sit on top of each other.
               'output'           : '_fhiaims.out',
               'runargs'          : '',
               'settings'         : '',
               'basis'            : '',
               # Allow path to be provided for the FHI-aims basis set
               'basis_dir'        : '',
               # The following allows partitioned fragments to be setup correctly. 
               # If element X is denoted by e.g. 'X1' in self.frag, setting qm_basis_alternative_labels = {'X': 1}
               # allows the specified basis for element X to be applied to species X1. 
               'qm_basis_alternative_labels'   : {},
               # Only works with separation of ecps (psps in FHI-aims language). AJL, 25/05/20
               # TWK 23/6/2020 Note separate ECPs are only applicable to the ionic embedding model, will be overriden for covalent coupling
               'separate_ecps'      : True,
               # These are all hardcoded options as they were in the previous tests
               # However, it makes more sense to me if we just take a dictionary of options and add all to the input, and remove these specifics
               # TWK: I would advise to keep them listed as-is here. This is so that (1) ChemShell can flag up unknown (and therefore possibly
               # incorrect) options to the user, and (2) if they are ever needed at the Fortran level, they will be converted automatically by DL_PY2F.
               'occupation_type'    : '', #'gaussian 0.01',
               'mixer'              : '', #'pulay',
               'n_max_pulay'        : '', #'10',
               'charge_mix_param'   : '', #'0.5',
               'sc_accuracy_rho'    : '', #'1E-5',
               'sc_accuracy_eev'    : '', #'1E-3',
               'sc_accuracy_etot'   : '', #'1E-6',
               'sc_accuracy_forces' : '', #'1E-4',
               'sc_iter_limit'      : '', #'100',
               'spin'               : 'none',
               # This defaults to atomic_zora scalar if unset in FHI-aims. I think I'd rather that was the long-term behaviour here too?
               'relativistic'       : 'none',
               #'compute_forces'     : '.true.',
               #'final_forces_cleaned' : '.false.',
             }

    _internals = {
                   '_ecprange'     : [],
                   '_ghost_list'   : [],
                   '_syscmd'       : 'aims.x', # default system command to run
                   '_eV2Ha'        : 0.03674932453 # Using FHI-aims internal standard for converting as eV -> Ha = *(1/27.2113845)
                 }

    _synons = {
              }

    def setOutKeys(self):
        '''Define {task:keyword} for searching in output file'''

        # 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 = {                           
                          #AJL, March 2019: This is preferable as in Ha, but clashes with SCF outputs so not easy
                          # to pick "correct" value without hacking the code.
                          # 'energy'   :{ 'Total energy                  ': [0,0,1,4,4,1,1,1.0]} 
                          # Alternative here is converting from eV back to Ha. Not pretty but works
                          # eV -> Ha = *(1/27.2113845) = *0.03674932453 (Ha value taken from FHI-aims output, not unitconvert)
                          # YL 07/03/2021: trying to correct the regex of literal |
                          'energy'   :{ '\\| Total energy of the DFT ': [0,0,1,11,11,1,1,self._eV2Ha]},
                          # Receive forces (-gradients) in eV/Angstrom. Converting as eV -> Ha = *(1/27.2113845) = *0.03674932453 (Ha value taken from FHI-aims output, not unitconvert)
                          # Converting as angstrom^-1 -> bohr^-1 = /unitconvert.Angstrom2Bohr 
                          'gradients':{ 'Total atomic forces': [1,1+self.frag.natoms-1,1,2,4,1,1,-(self._eV2Ha/unitconvert.Angstrom2Bohr)]},
                          'bq_gradients':{ 'Total external charge forces': [1,self.frag.nbqs-len(self._ecprange),1,2,4,1,1,-(self._eV2Ha/unitconvert.Angstrom2Bohr)]},
                          'empty_gradients': { 'Total atomic forces': [self.frag.natoms+1,self.frag.natoms+len(self._ghost_list),
                                                                       1,2,4,1,1,-(self._eV2Ha/unitconvert.Angstrom2Bohr)]},
                          'psp_gradients': { 'Total atomic forces': [self.frag.natoms+len(self._ghost_list)+1,self.frag.natoms+len(self._ghost_list)+len(self._ecprange),
                                                                       1,2,4,1,1,-(self._eV2Ha/unitconvert.Angstrom2Bohr)]},                           
                        }

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

        from ...data    import atomicdata
        from os.path    import join

        # Generate control file
        if input == self.input[0]:
            # Set the charge
            strbuff = '## ChemShell addition: charge ##' + '\n'
            strbuff += 'charge ' + str(self.charge) + '\n'
            strbuff += '## ChemShell addition: charge ##' + '\n'

            # Import settings
            if self.settings: 
                strbuff += utils.fileutils.file2text(self.settings, prefix=parallel.prefix) + '\n'

            # Calculator independent settings.
            if self.functional and 'xc ' not in strbuff:
                strbuff += 'xc ' + self._synons.get(self.functional, self.functional) + '\n'
            if self.spin and 'spin ' not in strbuff:
                strbuff += 'spin ' + self.spin + '\n'

            if 'qmmm' not in strbuff:
                strbuff += 'qmmm' '\n'

            # Add any user specific settings
            # I think these can be put into a dictionary, as they are not generalisable to other codes, and then we can just run a for loop through options
            if self.relativistic and 'relativistic ' not in strbuff:
                strbuff += 'relativistic ' + self.relativistic + '\n'
            if self.occupation_type and 'occupation_type' not in strbuff:
                strbuff += 'occupation_type ' + self.occupation_type + '\n'
            if self.mixer and 'mixer' not in strbuff:
                strbuff += 'mixer  ' + self.mixer + '\n'
            if self.n_max_pulay and 'n_max_pulay' not in strbuff: 
                strbuff += 'n_max_pulay  ' + self.n_max_pulay + '\n'
            if self.charge_mix_param and 'charge_mix_param' not in strbuff:
                strbuff += 'charge_mix_param  ' + self.charge_mix_param + '\n'
            if self.sc_accuracy_rho:
                strbuff += 'sc_accuracy_rho  ' + self.sc_accuracy_rho + '\n'
            if self.sc_accuracy_eev and 'sc_accurracy_eev' not in strbuff:
                strbuff += 'sc_accuracy_eev  ' + self.sc_accuracy_eev + '\n'
            if self.sc_accuracy_etot and 'sc_accurracy_etot' not in strbuff:
                strbuff += 'sc_accuracy_etot  ' + self.sc_accuracy_etot + '\n'
            if self.sc_accuracy_forces and 'sc_accurracy_forces' not in strbuff:
                strbuff += 'sc_accuracy_forces  ' + self.sc_accuracy_forces + '\n'
            if self.sc_iter_limit and 'sc_iter_limit' not in strbuff:
                strbuff += 'sc_iter_limit  ' + self.sc_iter_limit + '\n'

            # Enable processes necessary for calculating gradients
            if self._gradients:
                if 'compute_forces' not in strbuff:
                    strbuff += 'compute_forces .true.' + '\n'
                if 'final_forces_cleaned' not in strbuff:
                    strbuff += 'final_forces_cleaned .false.' + '\n'

            if self._mult >= 1 and self.spin == 'none':
                print('****')
                print('**** calculation settings for "spin" need to be set to "collinear". Please correct and retry')
                print('****')
                raise errors.ChemShGeneralError('If spin is greater than zero, an error will occur')

            

            # Import basis for QM atoms
            if self.basis:
                strbuff += utils.fileutils.file2text(self.basis, prefix=parallel.prefix) + '\n'
            elif self.basis_dir:
                # Extract elements
                znum_list = set(self.frag.znums)
                element_list = [str(atomicdata.symbols[z]) for z in znum_list]
                
                for i in range(len(element_list)):
                    # Following FHI-aims basis set file notation
                    file = join(self.basis_dir, f"{znum_list[i]:02}_{element_list[i]}_default")
                    for line in open(file, 'r'):
                        # This should only be used in the case of partitioned fragments where element X is denoted by e.g. 'X1' rather 
                        # than just 'X' in self.frag. For example, if self.qm_basis_alternative_labels = {'X': 1}, the specified basis
                        # set (from self.basis_dir) for element X will be pulled, with the line "species X" being replaced by 'species X1'
                        if element_list[i] in self.qm_basis_alternative_labels:
                            if 'species  ' in line:
                                line = line.rstrip() + str(int(self.qm_basis_alternative_labels[element_list[i]])) + '\n'
                        strbuff += line
            else:
                print('****')
                print('**** No calculation basis is detected. Please correct and retry')
                print('****')
                raise errors.ChemShGeneralError('Please set calculator basis= or basis_dir= for FHIaims')

            return strbuff
 
        # Generate geometry file
        else: # input == self.input[1]
            return self.getGeomBuff()

        # Catch any loose endings
        return

    def getGeomBuff(self):
        '''String buff containing the geometry information'''

        from numpy import __version__
        # TD 18/06/2024 : Numpy v2.0.0 version dependence
        from numpy.lib import NumpyVersion
        if NumpyVersion(__version__) >= '2.0.0b1':
            from numpy._core import defchararray, full
        else:
            from numpy.core import defchararray, full

        strbuff = '#FHI-aims ChemShell calculation'+ '\n'        
        # Print atoms
        strbuff_atoms = self.frag.coords2str(coords=9, atoms='names', order='ca', indent=' ', delimiter=' ',
                                             inserts=[(0,'atom')], unit='angstrom')

        if self.settings:
            pass
        else:
            if self.spin != 'none':
                for item in self.frag.names:
                    atom_found = False
                    for key in self.initial_moment:
                        if item.decode('UTF-8') == key:
                            atom_found = True
                            break
                    if atom_found == False:
                        raise errors.ChemShGeneralError('Not all atoms have been specified')

        # Enumerate as index i directly corresponds to atom number
        for i, line in enumerate(strbuff_atoms.splitlines()):
            line += '\n'
            if line.split()[4] in self.initial_moment:
                line += 'initial_moment ' + str(self.initial_moment[line.split()[4]]) + '\n'
            if i in self.initial_moment:
                line += 'initial_moment ' + str(self.initial_moment[i]) + '\n'    
            if line.split()[4] in self.initial_charge:
                line += 'initial_charge ' + str(self.initial_charge[line.split()[4]]) + '\n'
            if i in self.initial_charge:
                line += 'initial_charge ' + str(self.initial_charge[i]) + '\n'     
            strbuff += line
                                        
        # Print ECPs. Only occurs if doing QMMM, else self.ecp is undefined.
        if self.ecps:
            self._ecprange = self.frag.bqs.getRegion(2)
            # prefix with 'bq_' and suffix with '_e'
            ecpnames = defchararray.add(full((self.ecps.natoms,), 'bq_', dtype='S3'), self.ecps.names)
            ecpnames = defchararray.add(ecpnames, full((self.ecps.natoms,), '_e', dtype='S2'))
            strbuff += self.ecps.coords2str(coords=9, atoms=None, order='ca', indent=' ', delimiter=' ',
                                            inserts=[(0,'pseudocore'),(4,ecpnames)], unit='angstrom')

        # Print BQs
        if self.frag.nbqs:
            multipole_str = 'multipole'
            multipole_order = '0'
            atomstr = 'empty'
            self._ghost_list = []

            # nbqs is initialised to zero so no harm in just setting this for loop to run over this variable
            for ibq in range(self.frag.nbqs):

                # Check BQ is not an ECP centre. If no ECPs, this list is empty
                if ibq not in self._ecprange:
                    # AJL, March 2019: Get BQ coords, convert to Angstrom
                    # Made this an if statement so I can match settings for the TCL version of ChemShell (giving Bohr2Angstrom) and validate setup
                    coords = [ (float(coord)*unitconvert.Bohr2Angstrom) for coord in list(self.frag.bqs.coords[ibq]) ]
                    strbuff = ''.join([strbuff, '%10s    %16.9f    %16.9f    %16.9f   %2s   %12.8f\n'%(multipole_str, coords[0], coords[1], coords[2],
                                       multipole_order, self.frag.bqs.charges[ibq])])

                    # AJL, April 2019: Check to see if we need an empty site on this BQ. 
                    # In short, this is necessary as it removes the integration grids from being right on top of the point charge during the
                    # numerical integrations. Without it, you get spurious, uncontrollable results with singularities if the BQs are close to QM atoms
                    # The default ghost_cutoff is set to 5 angstrom, so by default this check is made. First we look at atom-bq interactions
                    if self.frag.selectByRadius(radius=self.ghost_max_cutoff*unitconvert.Angstrom2Bohr,centre=list(self.frag.bqs.coords[ibq])).size > 0:

                        if self.ghost_species:
                            # Now we need to check for bq-bq interactions. We cannot put empty sites on top of each other so if we have
                            # BQs very close together (0.001 au) then we only add an empty site over the first instance
                            ghost_close_neighbours = self.frag.bqs.selectByRadius(radius=self.ghost_min_cutoff*unitconvert.Angstrom2Bohr, centre=list(self.frag.bqs.coords[ibq]))

                            if len(ghost_close_neighbours) == 1 or any(jbq > ibq for jbq in ghost_close_neighbours):
                                #print('**** Adding an integration grid to protect from QM-BQ energy singularities')
                                strbuff = ''.join([strbuff, '%10s    %16.9f    %16.9f    %16.9f   %7s\n'%(atomstr, coords[0], coords[1], coords[2],
                                                   self.ghost_species)])
                                #Need to keep count of the empty sites we add so that we can include gradients information post-calculation
                                self._ghost_list.append(ibq)

                        else:
                            # I've made this verbose so the user understands the error
                            print('****')
                            print('**** A QM and BQ site have been identified in close proximity. Due to the numerical')
                            print('**** integrations used in FHI-aims, an "empty" site **must** be placed on the BQ to protect')
                            print('**** against Coulomb singularities. Please define the "empty" site species (ghost_species)')
                            print('**** Note: this error can be silenced by setting ghost_max_cutoff = 0.')
                            print('****')
                            raise errors.ChemShGeneralError('ghost_species must have a species name for the empty sites. Check control.in and update input')

        return strbuff

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

        from numpy import insert as np_insert
        from numpy import zeros  as np_zeros
        from ...io import file

        self.setOutKeys()
        self._result.energy    = file.getArrayFromFile(self.output, 'out', self._result.energy.shape   , regex=self._outkeys['energy'])
        if self._result.energy == 0.0:
            raise errors.ChemShGeneralError('Energy seems to be an invalid number. Check FHI-aims output.')

        try:
            if self._gradients:
                self._result.gradients = file.getArrayFromFile(self.output, 'out', self._result.gradients.shape   , regex=self._outkeys['gradients'])
        except:
            print("\n >>> ChemShell FHI-aims: failed parsing gradients from output:\n")
            raise

        # AJL 29/05/20: Tried to make this more akin to NWChem interface, i.e. ECP gradients are read into a separate object from BQ gradients
        if self._gradients and self.ecps:
            self.ecps.gradients = file.getArrayFromFile(self.output, 'out', (len(self._ecprange),3), regex=self._outkeys['psp_gradients'])

        if self._gradients and self.frag.nbqs:
            gradients_bq_sites = file.getArrayFromFile(self.output, 'out', (self.frag.nbqs-len(self._ecprange),3), regex=self._outkeys['bq_gradients'])

            # When we use PSPs (ECPs by another name) in FHI-aims, we do not also include the BQ centres in the input geometry. This means that the
            # structure for bqs.gradients must be recreated by adding in some zeros at the indices for PSP centres.
            # Previously, we just inserted the PSP gradients here, but in line with the recombination of ecps.gradients and bq.gradients in qm.py, this has been changed. AJL 29/05/20
            if self.ecps:
                #gradients_bq_sites = np_insert(gradients_bq_sites,self._ecprange,gradients_psp_sites,axis=0)
                for ipsp in range(len(self._ecprange)):
                    #print(type(self.ecps.gradients[ipsp]), self.ecps.gradients[ipsp].shape, self.ecps.gradients[ipsp])
                    gradients_bq_sites = np_insert(gradients_bq_sites,self._ecprange[ipsp],np_zeros(3),axis=0)
                #print(self.frag.bqs.gradients.shape, g]radients_bq_sites.shape)

            if len(self.frag.bqs.gradients) != len(gradients_bq_sites):
                raise errors.ChemShGeneralError('Array length mismatch when reading gradients for BQs and pseudopotentials. Check FHI-aims output.')
            self.frag.bqs.gradients = gradients_bq_sites
                
            # If we have empty sites, we need to include the gradients on these with the gradients on the BQs as well.
            if any(self._ghost_list):
                gradients_empty_sites = file.getArrayFromFile(self.output, 'out', (len(self._ghost_list),3), regex=self._outkeys['empty_gradients'])

                for ighost in range(len(self._ghost_list)):
                    self.frag.bqs.gradients[self._ghost_list[ighost]] += gradients_empty_sites[ighost]


    def runLib(self, **kwargs):
        '''Run a library call of FHI-aims (called by run.runLib())'''

        from ctypes     import byref 
        from ...dl_py2f import py2f

        print('\n >>> Directly-linked FHI-aims library will be used...')

        libfhiaims = utils.modutils.getLinkedLib(self.__class__.__name__, chemsh.CHEMSH_LIB_DIR)

        libfhiaims.interface_fhiaims(byref(py2f(self.frag)), byref(py2f(self.options)))

        return
