# Source code for qcdb.interface_gcp

    #
    # @BEGIN LICENSE
    #
    # Psi4: an open-source quantum chemistry software package
    #
    # Copyright (c) 2007-2019 The Psi4 Developers.
    #
    # The copyrights for code used from other parties are included in
    # the corresponding files.
    #
    # This file is part of Psi4.
    #
    # Psi4 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, version 3.
    #
    # Psi4 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 Psi4; if not, write to the Free Software Foundation, Inc.,
    # 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
    #
    # @END LICENSE
    #

    """Module with functions that interface with Grimme's GCP code."""
    import os
    import re
    import uuid
    import shutil
    import socket
    import subprocess

    try:
        from psi4.driver.p4util.exceptions import *
        from psi4 import core
        isP4regime = True
    except ImportError:
        from .exceptions import *
        isP4regime = False
    from .util import parse_dertype
    from .molecule import Molecule

    def run_gcp(self, func=None, dertype=None, verbose=False):  # dashlvl=None, dashparam=None
        """Function to call Grimme's GCP program
        https://www.chemie.uni-bonn.de/pctc/mulliken-center/software/gcp/gcp
        to compute an a posteriori geometrical BSSE correction to *self* for
        several HF, generic DFT, and specific HF-3c and PBEh-3c method/basis
        combinations, *func*. Returns energy if *dertype* is 0, gradient
        if *dertype* is 1, else tuple of energy and gradient if *dertype*
        unspecified. The gcp executable must be independently compiled and
        found in :envvar:`PATH` or :envvar:`PSIPATH`. *self* may be either a
        qcdb.Molecule (sensibly) or a psi4.Molecule (works b/c psi4.Molecule
        has been extended by this method py-side and only public interface
        fns used) or a string that can be instantiated into a qcdb.Molecule.

        """
        # Create (if necessary) and update qcdb.Molecule
        if isinstance(self, Molecule):
            # called on a qcdb.Molecule
            pass
        elif isinstance(self, core.Molecule):
            # called on a python export of a psi4.core.Molecule (py-side through Psi4's driver)
            self.create_psi4_string_from_molecule()
        elif isinstance(self, str):
            # called on a string representation of a psi4.Molecule (c-side through psi4.Dispersion)
            self = Molecule(self)
        else:
            raise ValidationError("""Argument mol must be psi4string or qcdb.Molecule""")
        self.update_geometry()

    #    # Validate arguments
    #    dashlvl = dashlvl.lower()
    #    dashlvl = dash_alias['-' + dashlvl][1:] if ('-' + dashlvl) in dash_alias.keys() else dashlvl
    #    if dashlvl not in dashcoeff.keys():
    #        raise ValidationError("""-D correction level %s is not available. Choose among %s.""" % (dashlvl, dashcoeff.keys()))

        if dertype is None:
            derint, derdriver = -1, 'gradient'
        else:
            derint, derdriver = parse_dertype(dertype, max_derivative=1)

    #    if func is None:
    #        if dashparam is None:
    #            # defunct case
    #            raise ValidationError("""Parameters for -D correction missing. Provide a func or a dashparam kwarg.""")
    #        else:
    #            # case where all param read from dashparam dict (which must have all correct keys)
    #            func = 'custom'
    #            dashcoeff[dashlvl][func] = {}
    #            dashparam = dict((k.lower(), v) for k, v in dashparam.iteritems())
    #            for key in dashcoeff[dashlvl]['b3lyp'].keys():
    #                if key in dashparam.keys():
    #                    dashcoeff[dashlvl][func][key] = dashparam[key]
    #                else:
    #                    raise ValidationError("""Parameter %s is missing from dashparam dict %s.""" % (key, dashparam))
    #    else:
    #        func = func.lower()
    #        if func not in dashcoeff[dashlvl].keys():
    #            raise ValidationError("""Functional %s is not available for -D level %s.""" % (func, dashlvl))
    #        if dashparam is None:
    #            # (normal) case where all param taken from dashcoeff above
    #            pass
    #        else:
    #            # case where items in dashparam dict can override param taken from dashcoeff above
    #            dashparam = dict((k.lower(), v) for k, v in dashparam.iteritems())
    #            for key in dashcoeff[dashlvl]['b3lyp'].keys():
    #                if key in dashparam.keys():
    #                    dashcoeff[dashlvl][func][key] = dashparam[key]

        # TODO temp until figure out paramfile
        allowed_funcs = ['HF/MINIS', 'DFT/MINIS', 'HF/MINIX', 'DFT/MINIX',
            'HF/SV', 'DFT/SV', 'HF/def2-SV(P)', 'DFT/def2-SV(P)', 'HF/def2-SVP',
            'DFT/def2-SVP', 'HF/DZP', 'DFT/DZP', 'HF/def-TZVP', 'DFT/def-TZVP',
            'HF/def2-TZVP', 'DFT/def2-TZVP', 'HF/631Gd', 'DFT/631Gd',
            'HF/def2-TZVP', 'DFT/def2-TZVP', 'HF/cc-pVDZ', 'DFT/cc-pVDZ',
            'HF/aug-cc-pVDZ', 'DFT/aug-cc-pVDZ', 'DFT/SV(P/h,c)', 'DFT/LANL',
            'DFT/pobTZVP', 'TPSS/def2-SVP', 'PW6B95/def2-SVP',
            # specials
            'hf3c', 'pbeh3c']
        allowed_funcs = [f.lower() for f in allowed_funcs]
        if func.lower() not in allowed_funcs:
            raise Dftd3Error("""bad gCP func: %s. need one of: %r""" % (func, allowed_funcs))

        # Move ~/.dftd3par.<hostname> out of the way so it won't interfere
        defaultfile = os.path.expanduser('~') + '/.dftd3par.' + socket.gethostname()
        defmoved = False
        if os.path.isfile(defaultfile):
            os.rename(defaultfile, defaultfile + '_hide')
            defmoved = True

        # Find environment by merging PSIPATH and PATH environment variables
        lenv = {
            'PATH': ':'.join([os.path.abspath(x) for x in os.environ.get('PSIPATH', '').split(':') if x != '']) + \
                    ':' + os.environ.get('PATH'),
            'LD_LIBRARY_PATH': os.environ.get('LD_LIBRARY_PATH')
            }
        #   Filter out None values as subprocess will fault on them
        lenv = {k: v for k, v in lenv.items() if v is not None}

        # Find out if running from Psi4 for scratch details and such
        try:
            import psi4
        except ImportError as err:
            isP4regime = False
        else:
            isP4regime = True

        # Setup unique scratch directory and move in
        current_directory = os.getcwd()
        if isP4regime:
            psioh = core.IOManager.shared_object()
            psio = core.IO.shared_object()
            os.chdir(psioh.get_default_path())
            gcp_tmpdir = 'psi.' + str(os.getpid()) + '.' + psio.get_default_namespace() + \
                '.gcp.' + str(uuid.uuid4())[:8]
        else:
            gcp_tmpdir = os.environ['HOME'] + os.sep + 'gcp_' + str(uuid.uuid4())[:8]
        if os.path.exists(gcp_tmpdir) is False:
            os.mkdir(gcp_tmpdir)
        os.chdir(gcp_tmpdir)

        # Write gcp_parameters file that governs cp correction
    #    paramcontents = gcp_server(func, dashlvl, 'dftd3')
    #    paramfile1 = 'dftd3_parameters'  # older patched name
    #    with open(paramfile1, 'w') as handle:
    #        handle.write(paramcontents)
    #    paramfile2 = '.gcppar'
    #    with open(paramfile2, 'w') as handle:
    #        handle.write(paramcontents)

    ###Two kinds of parameter files can be read in: A short and an extended version. Both are read from
    ###$HOME/.gcppar.$HOSTNAME by default. If the option -local is specified the file is read in from
    ###the current working directory: .gcppar
    ###The short version reads in: basis-keywo

        # Write dftd3_geometry file that supplies geometry to dispersion calc
        numAtoms = self.natom()
        geom = self.save_string_xyz()
        reals = []
        for line in geom.splitlines():
            lline = line.split()
            if len(lline) != 4:
                continue
            if lline[0] == 'Gh':
                numAtoms -= 1
            else:
                reals.append(line)

        geomtext = str(numAtoms) + '\n\n'
        for line in reals:
            geomtext += line.strip() + '\n'
        geomfile = './gcp_geometry.xyz'
        with open(geomfile, 'w') as handle:
            handle.write(geomtext)
        # TODO somehow the variations on save_string_xyz and
        #   whether natom and chgmult does or doesn't get written
        #   have gotten all tangled. I fear this doesn't work
        #   the same btwn libmints and qcdb or for ghosts

        # Call gcp program
        command = ['gcp', geomfile]
        command.extend(['-level', func])
        if derint != 0:
            command.append('-grad')
        try:
            #print('command', command)
            dashout = subprocess.Popen(command, stdout=subprocess.PIPE, env=lenv)
        except OSError as e:
            raise ValidationError('Program gcp not found in path. %s' % e)
        out, err = dashout.communicate()

        # Parse output
        success = False
        for line in out.splitlines():
            line = line.decode('utf-8')
            if re.match('  Egcp:', line):
                sline = line.split()
                dashd = float(sline[1])
            if re.match('     normal termination of gCP', line):
                success = True

        if not success:
            os.chdir(current_directory)
            raise Dftd3Error("""Unsuccessful gCP run.""")

        # Parse grad output
        if derint != 0:
            derivfile = './gcp_gradient'
            dfile = open(derivfile, 'r')
            dashdderiv = []
            for line in geom.splitlines():
                lline = line.split()
                if len(lline) != 4:
                    continue
                if lline[0] == 'Gh':
                    dashdderiv.append([0.0, 0.0, 0.0])
                else:
                    dashdderiv.append([float(x.replace('D', 'E')) for x in dfile.readline().split()])
            dfile.close()

            if len(dashdderiv) != self.natom():
                raise ValidationError('Program gcp gradient file has %d atoms- %d expected.' % \
                    (len(dashdderiv), self.natom()))

        # Prepare results for Psi4
        if isP4regime and derint != 0:
            core.set_variable('GCP CORRECTION ENERGY', dashd)
            psi_dashdderiv = core.Matrix.from_list(dashdderiv)

        # Print program output to file if verbose
        if not verbose and isP4regime:
            verbose = True if core.get_option('SCF', 'PRINT') >= 3 else False
        if verbose:

            text = '\n  ==> GCP Output <==\n'
            text += out.decode('utf-8')
            if derint != 0:
                with open(derivfile, 'r') as handle:
                    text += handle.read().replace('D', 'E')
                text += '\n'
            if isP4regime:
                core.print_out(text)
            else:
                print(text)

    #    # Clean up files and remove scratch directory
    #    os.unlink(paramfile1)
    #    os.unlink(paramfile2)
    #    os.unlink(geomfile)
    #    if derint != 0:
    #        os.unlink(derivfile)
    #    if defmoved is True:
    #        os.rename(defaultfile + '_hide', defaultfile)

        # clean up files and remove scratch directory
        os.chdir('..')
        try:
            shutil.rmtree(gcp_tmpdir)
        except OSError as err:
            raise OSError('Unable to remove gcp temporary directory: {}'.format(gcp_tmpdir)) from err
        os.chdir(current_directory)

        # return -D & d(-D)/dx
        if derint == -1:
            return dashd, dashdderiv
        elif derint == 0:
            return dashd
        elif derint == 1:
            return psi_dashdderiv

    try:
        # Attach method to libmints psi4.Molecule class
        core.Molecule.run_gcp = run_gcp
    except (NameError, AttributeError):
        # But don't worry if that doesn't work b/c
        #   it'll get attached to qcdb.Molecule class
        pass