# Source code for psi4.driver.p4util.fcidump

    #
    # @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 utility function for dumping the Hamiltonian to file in FCIDUMP format."""

    from datetime import datetime

    import numpy as np

    from psi4.driver import psifiles as psif
    from psi4.driver.p4util.util import compare_values, success
    from psi4.driver.procrouting.proc_util import check_iwl_file_from_scf_type

    from psi4 import core
    from .exceptions import ValidationError, TestComparisonError

    def fcidump(wfn, fname='INTDUMP', oe_ints=None):
        """Save integrals to file in FCIDUMP format as defined in Comp. Phys. Commun. 54 75 (1989)
        Additional one-electron integrals, including orbital energies, can also be saved.
        This latter format can be used with the HANDE QMC code but is not standard.

        :returns: None

        :raises: ValidationError when SCF wavefunction is not RHF

        :type wfn: :py:class:`~psi4.core.Wavefunction`
        :param wfn: set of molecule, basis, orbitals from which to generate cube files
        :param fname: name of the integrals file, defaults to INTDUMP
        :param oe_ints: list of additional one-electron integrals to save to file.
        So far only EIGENVALUES is a valid option.

        :examples:

        >>> # [1] Save one- and two-electron integrals to standard FCIDUMP format
        >>> E, wfn = energy('scf', return_wfn=True)
        >>> fcidump(wfn)

        >>> # [2] Save orbital energies, one- and two-electron integrals.
        >>> E, wfn = energy('scf', return_wfn=True)
        >>> fcidump(wfn, oe_ints=['EIGENVALUES'])

        """
        # Get some options
        reference = core.get_option('SCF', 'REFERENCE')
        ints_tolerance = core.get_global_option('INTS_TOLERANCE')
        # Some sanity checks
        if reference not in ['RHF', 'UHF']:
            raise ValidationError('FCIDUMP not implemented for {} references\n'.format(reference))
        if oe_ints is None:
            oe_ints = []

        molecule = wfn.molecule()
        docc = wfn.doccpi()
        frzcpi = wfn.frzcpi()
        frzvpi = wfn.frzvpi()
        active_docc = docc - frzcpi
        active_socc = wfn.soccpi()
        active_mopi = wfn.nmopi() - frzcpi - frzvpi

        nbf = active_mopi.sum() if wfn.same_a_b_orbs() else 2 * active_mopi.sum()
        nirrep = wfn.nirrep()
        nelectron = 2 * active_docc.sum() + active_socc.sum()
        irrep_map = _irrep_map(wfn)

        wfn_irrep = 0
        for h, n_socc in enumerate(active_socc):
            if n_socc % 2 == 1:
                wfn_irrep ^= h

        core.print_out('Writing integrals in FCIDUMP format to ' + fname + '\n')
        # Generate FCIDUMP header
        header = '&FCI\n'
        header += 'NORB={:d},\n'.format(nbf)
        header += 'NELEC={:d},\n'.format(nelectron)
        header += 'MS2={:d},\n'.format(wfn.nalpha() - wfn.nbeta())
        header += 'UHF=.{}.,\n'.format(not wfn.same_a_b_orbs()).upper()
        orbsym = ''
        for h in range(active_mopi.n()):
            for n in range(frzcpi[h], frzcpi[h] + active_mopi[h]):
                orbsym += '{:d},'.format(irrep_map[h])
                if not wfn.same_a_b_orbs():
                    orbsym += '{:d},'.format(irrep_map[h])
        header += 'ORBSYM={}\n'.format(orbsym)
        header += 'ISYM={:d},\n'.format(irrep_map[wfn_irrep])
        header += '&END\n'
        with open(fname, 'w') as intdump:
            intdump.write(header)

        # Get an IntegralTransform object
        check_iwl_file_from_scf_type(core.get_global_option('SCF_TYPE'), wfn)
        spaces = [core.MOSpace.all()]
        trans_type = core.IntegralTransform.TransformationType.Restricted
        if not wfn.same_a_b_orbs():
            trans_type = core.IntegralTransform.TransformationType.Unrestricted
        ints = core.IntegralTransform(wfn, spaces, trans_type)
        ints.transform_tei(core.MOSpace.all(), core.MOSpace.all(), core.MOSpace.all(), core.MOSpace.all())
        core.print_out('Integral transformation complete!\n')

        DPD_info = {'instance_id': ints.get_dpd_id(), 'alpha_MO': ints.DPD_ID('[A>=A]+'), 'beta_MO': 0}
        if not wfn.same_a_b_orbs():
            DPD_info['beta_MO'] = ints.DPD_ID("[a>=a]+")
        # Write TEI to fname in FCIDUMP format
        core.fcidump_tei_helper(nirrep, wfn.same_a_b_orbs(), DPD_info, ints_tolerance, fname)

        # Read-in OEI and write them to fname in FCIDUMP format
        # Indexing functions to translate from zero-based (C and Python) to
        # one-based (Fortran)
        mo_idx = lambda x: x + 1
        alpha_mo_idx = lambda x: 2 * x + 1
        beta_mo_idx = lambda x: 2 * (x + 1)

        with open(fname, 'a') as intdump:
            core.print_out('Writing frozen core operator in FCIDUMP format to ' + fname + '\n')
            if reference == 'RHF':
                PSIF_MO_FZC = 'MO-basis Frozen-Core Operator'
                moH = core.Matrix(PSIF_MO_FZC, wfn.nmopi(), wfn.nmopi())
                moH.load(core.IO.shared_object(), psif.PSIF_OEI)
                mo_slice = core.Slice(frzcpi, active_mopi)
                MO_FZC = moH.get_block(mo_slice, mo_slice)
                offset = 0
                for h, block in enumerate(MO_FZC.nph):
                    il = np.tril_indices(block.shape[0])
                    for index, x in np.ndenumerate(block[il]):
                        row = mo_idx(il[0][index] + offset)
                        col = mo_idx(il[1][index] + offset)
                        if (abs(x) > ints_tolerance):
                            intdump.write('{:29.20E} {:4d} {:4d} {:4d} {:4d}\n'.format(x, row, col, 0, 0))
                    offset += block.shape[0]
                # Additional one-electron integrals as requested in oe_ints
                # Orbital energies
                core.print_out('Writing orbital energies in FCIDUMP format to ' + fname + '\n')
                if 'EIGENVALUES' in oe_ints:
                    eigs_dump = write_eigenvalues(wfn.epsilon_a().get_block(mo_slice).to_array(), mo_idx)
                    intdump.write(eigs_dump)
            else:
                PSIF_MO_A_FZC = 'MO-basis Alpha Frozen-Core Oper'
                moH_A = core.Matrix(PSIF_MO_A_FZC, wfn.nmopi(), wfn.nmopi())
                moH_A.load(core.IO.shared_object(), psif.PSIF_OEI)
                mo_slice = core.Slice(frzcpi, active_mopi)
                MO_FZC_A = moH_A.get_block(mo_slice, mo_slice)
                offset = 0
                for h, block in enumerate(MO_FZC_A.nph):
                    il = np.tril_indices(block.shape[0])
                    for index, x in np.ndenumerate(block[il]):
                        row = alpha_mo_idx(il[0][index] + offset)
                        col = alpha_mo_idx(il[1][index] + offset)
                        if (abs(x) > ints_tolerance):
                            intdump.write('{:29.20E} {:4d} {:4d} {:4d} {:4d}\n'.format(x, row, col, 0, 0))
                    offset += block.shape[0]
                PSIF_MO_B_FZC = 'MO-basis Beta Frozen-Core Oper'
                moH_B = core.Matrix(PSIF_MO_B_FZC, wfn.nmopi(), wfn.nmopi())
                moH_B.load(core.IO.shared_object(), psif.PSIF_OEI)
                mo_slice = core.Slice(frzcpi, active_mopi)
                MO_FZC_B = moH_B.get_block(mo_slice, mo_slice)
                offset = 0
                for h, block in enumerate(MO_FZC_B.nph):
                    il = np.tril_indices(block.shape[0])
                    for index, x in np.ndenumerate(block[il]):
                        row = beta_mo_idx(il[0][index] + offset)
                        col = beta_mo_idx(il[1][index] + offset)
                        if (abs(x) > ints_tolerance):
                            intdump.write('{:29.20E} {:4d} {:4d} {:4d} {:4d}\n'.format(x, row, col, 0, 0))
                    offset += block.shape[0]
                # Additional one-electron integrals as requested in oe_ints
                # Orbital energies
                core.print_out('Writing orbital energies in FCIDUMP format to ' + fname + '\n')
                if 'EIGENVALUES' in oe_ints:
                    alpha_eigs_dump = write_eigenvalues(wfn.epsilon_a().get_block(mo_slice).to_array(), alpha_mo_idx)
                    beta_eigs_dump = write_eigenvalues(wfn.epsilon_b().get_block(mo_slice).to_array(), beta_mo_idx)
                    intdump.write(alpha_eigs_dump + beta_eigs_dump)
            # Dipole integrals
            #core.print_out('Writing dipole moment OEI in FCIDUMP format to ' + fname + '\n')
            # Traceless quadrupole integrals
            #core.print_out('Writing traceless quadrupole moment OEI in FCIDUMP format to ' + fname + '\n')
            # Frozen core + nuclear repulsion energy
            core.print_out('Writing frozen core + nuclear repulsion energy in FCIDUMP format to ' + fname + '\n')
            e_fzc = ints.get_frozen_core_energy()
            e_nuc = molecule.nuclear_repulsion_energy(wfn.get_dipole_field_strength())
            intdump.write('{: 29.20E} {:4d} {:4d} {:4d} {:4d}\n'.format(e_fzc + e_nuc, 0, 0, 0, 0))
        core.print_out('Done generating {} with integrals in FCIDUMP format.\n'.format(fname))

    def write_eigenvalues(eigs, mo_idx):
        """Prepare multi-line string with one-particle eigenvalues to be written to the FCIDUMP file.
        """
        eigs_dump = ''
        iorb = 0
        for h, block in enumerate(eigs):
            for idx, x in np.ndenumerate(block):
                eigs_dump += '{: 29.20E} {:4d} {:4d} {:4d} {:4d}\n'.format(x, mo_idx(iorb), 0, 0, 0)
                iorb += 1
        return eigs_dump

    def _irrep_map(wfn):
        """Returns an array of irrep indices that maps from Psi4's ordering convention to the standard FCIDUMP convention.
        """
        symm = wfn.molecule().point_group().symbol()
        psi2dump = {'c1' : [1],               # A
                    'ci' : [1,2],             # Ag Au
                    'c2' : [1,2],             # A  B
                    'cs' : [1,2],             # A' A"
                    'd2' : [1,4,3,2],         # A  B1  B2  B3
                    'c2v' : [1,4,2,3],        # A1 A2  B1  B2
                    'c2h' : [1,4,2,3],        # Ag Bg  Au  Bu
                    'd2h' : [1,4,6,7,8,5,3,2] # Ag B1g B2g B3g Au B1u B2u B3u
                    }

        irrep_map = psi2dump[symm]
        return np.array(irrep_map, dtype='int')

    def fcidump_from_file(fname):
        """Function to read in a FCIDUMP file.

        :returns: a dictionary with FCIDUMP header and integrals
        The key-value pairs are:
          - 'norb' : number of basis functions
          - 'nelec' : number of electrons
          - 'ms2' : spin polarization of the system
          - 'isym' : symmetry of state (if present in FCIDUMP)
          - 'orbsym' : list of symmetry labels of each orbital
          - 'uhf' : whether restricted or unrestricted
          - 'enuc' : nuclear repulsion plus frozen core energy
          - 'epsilon' : orbital energies
          - 'hcore' : core Hamiltonian
          - 'eri' : electron-repulsion integrals

        :param fname: FCIDUMP file name
        """
        intdump = {}
        with open(fname, 'r') as handle:
            assert '&FCI' == handle.readline().strip()

            skiplines = 1
            read = True
            while True:
                skiplines += 1
                line = handle.readline()
                if 'END' in line:
                    break

                key, value = line.split('=')
                value = value.strip().rstrip(',')
                if key == 'UHF':
                    value = 'TRUE' in value
                elif key == 'ORBSYM':
                    value = [int(x) for x in value.split(',')]
                else:
                    value = int(value.replace(',', ''))

                intdump[key.lower()] = value

        # Read the data and index, skip header
        raw_ints = np.genfromtxt(fname, skip_header=skiplines)

        # Read last line, i.e. Enuc + Efzc
        intdump['enuc'] = raw_ints[-1, 0]

        # Read in integrals and indices
        ints = raw_ints[:-1, 0]

        # Get dimensions and indices
        nbf = intdump['norb']
        idxs = raw_ints[:, 1:].astype(np.int) - 1

        # Slices
        sl = slice(ints.shape[0] - nbf, ints.shape[0])

        # Extract orbital energies
        epsilon = np.zeros(nbf)
        epsilon[idxs[sl, 0]] = ints[sl]
        intdump['epsilon'] = epsilon

        # Count how many 2-index intdump we have
        sl = slice(sl.start - nbf * nbf, sl.stop - nbf)
        two_index = np.all(idxs[sl, 2:] == -1, axis=1).sum()
        sl = slice(sl.stop - two_index, sl.stop)

        # Extract Hcore
        Hcore = np.zeros((nbf, nbf))
        Hcore[(idxs[sl, 0], idxs[sl, 1])] = ints[sl]
        Hcore[(idxs[sl, 1], idxs[sl, 0])] = ints[sl]
        intdump['hcore'] = Hcore

        # Extract ERIs
        sl = slice(0, sl.start)
        eri = np.zeros((nbf, nbf, nbf, nbf))
        eri[(idxs[sl, 0], idxs[sl, 1], idxs[sl, 2], idxs[sl, 3])] = ints[sl]
        eri[(idxs[sl, 0], idxs[sl, 1], idxs[sl, 3], idxs[sl, 2])] = ints[sl]
        eri[(idxs[sl, 1], idxs[sl, 0], idxs[sl, 2], idxs[sl, 3])] = ints[sl]
        eri[(idxs[sl, 1], idxs[sl, 0], idxs[sl, 3], idxs[sl, 2])] = ints[sl]
        eri[(idxs[sl, 2], idxs[sl, 3], idxs[sl, 0], idxs[sl, 1])] = ints[sl]
        eri[(idxs[sl, 3], idxs[sl, 2], idxs[sl, 0], idxs[sl, 1])] = ints[sl]
        eri[(idxs[sl, 2], idxs[sl, 3], idxs[sl, 1], idxs[sl, 0])] = ints[sl]
        eri[(idxs[sl, 3], idxs[sl, 2], idxs[sl, 1], idxs[sl, 0])] = ints[sl]
        intdump['eri'] = eri

        return intdump

    def compare_fcidumps(expected, computed, label):
        """Function to compare two FCIDUMP files. Prints :py:func:`util.success`
        when value *computed* matches value *expected*.
        Performs a system exit on failure. Used in input files in the test suite.

        :returns: a dictionary of energies computed from the MO integrals.
        The key-value pairs are:
          - 'NUCLEAR REPULSION ENERGY' : nuclear repulsion plus frozen core energy
          - 'ONE-ELECTRON ENERGY' : SCF one-electron energy
          - 'TWO-ELECTRON ENERGY' : SCF two-electron energy
          - 'SCF TOTAL ENERGY' : SCF total energy
          - 'MP2 CORRELATION ENERGY' : MP2 correlation energy

        :param expected: reference FCIDUMP file
        :param computed: computed FCIDUMP file
        :param label: string labelling the test
        """

        try:
            from deepdiff import DeepDiff
        except ImportError:
            raise ImportError(
                """Python module deepdiff not found. Solve by installing it: `conda install deepdiff -c conda-forge` or `pip install deepdiff`"""
            )

        # Grab expected header and integrals
        ref_intdump = fcidump_from_file(expected)
        intdump = fcidump_from_file(computed)

        # Compare headers
        header_diff = DeepDiff(
            ref_intdump,
            intdump,
            ignore_order=True,
            exclude_paths={"root['enuc']", "root['hcore']", "root['eri']", "root['epsilon']"})
        if header_diff:
            message = ("\tComputed FCIDUMP file header does not match expected header.\n")
            raise TestComparisonError(header_diff)

        ref_energies = energies_from_fcidump(ref_intdump)
        energies = energies_from_fcidump(intdump)

        pass_1el = compare_values(ref_energies['ONE-ELECTRON ENERGY'], energies['ONE-ELECTRON ENERGY'], 7,
                                  label + '. 1-electron energy')
        pass_2el = compare_values(ref_energies['TWO-ELECTRON ENERGY'], energies['TWO-ELECTRON ENERGY'], 7,
                                  label + '. 2-electron energy')
        pass_scf = compare_values(ref_energies['SCF TOTAL ENERGY'], energies['SCF TOTAL ENERGY'], 10,
                                  label + '. SCF total energy')
        pass_mp2 = compare_values(ref_energies['MP2 CORRELATION ENERGY'], energies['MP2 CORRELATION ENERGY'], 10,
                                  label + '. MP2 correlation energy')

        if (pass_1el and pass_2el and pass_scf and pass_mp2):
            success(label)

        return True

    def energies_from_fcidump(intdump):
        energies = {}
        energies['NUCLEAR REPULSION ENERGY'] = intdump['enuc']
        epsilon = intdump['epsilon']
        Hcore = intdump['hcore']
        eri = intdump['eri']

        # Compute SCF energy
        energies['ONE-ELECTRON ENERGY'], energies['TWO-ELECTRON ENERGY'] = _scf_energy(Hcore, eri,
                                                                                       np.where(epsilon < 0)[0],
                                                                                       intdump['uhf'])
        # yapf: disable
        energies['SCF TOTAL ENERGY'] = energies['ONE-ELECTRON ENERGY'] + energies['TWO-ELECTRON ENERGY'] + energies['NUCLEAR REPULSION ENERGY']
        # yapf: enable

        # Compute MP2 energy
        energies['MP2 CORRELATION ENERGY'] = _mp2_energy(eri, epsilon, intdump['uhf'])

        return energies

    def _scf_energy(Hcore, ERI, occ_sl, unrestricted):
        scf_1el_e = np.einsum('ii->', Hcore[np.ix_(occ_sl, occ_sl)])
        if not unrestricted:
            scf_1el_e *= 2
        coulomb = np.einsum('iijj->', ERI[np.ix_(occ_sl, occ_sl, occ_sl, occ_sl)])
        exchange = np.einsum('ijij->', ERI[np.ix_(occ_sl, occ_sl, occ_sl, occ_sl)])
        if unrestricted:
            scf_2el_e = 0.5 * (coulomb - exchange)
        else:
            scf_2el_e = 2.0 * coulomb - exchange

        return scf_1el_e, scf_2el_e

    def _mp2_energy(ERI, epsilon, unrestricted):
        # Occupied and virtual slices
        occ_sl = np.where(epsilon < 0)[0]
        vir_sl = np.where(epsilon > 0)[0]
        eocc = epsilon[occ_sl]
        evir = epsilon[vir_sl]
        denom = 1 / (eocc.reshape(-1, 1, 1, 1) - evir.reshape(-1, 1, 1) + eocc.reshape(-1, 1) - evir)
        MO = ERI[np.ix_(occ_sl, vir_sl, occ_sl, vir_sl)]
        if unrestricted:
            mp2_e = 0.5 * np.einsum("abrs,abrs,abrs->", MO, MO - MO.swapaxes(1, 3), denom)
        else:
            mp2_e = np.einsum('iajb,iajb,iajb->', MO, MO, denom) + np.einsum('iajb,iajb,iajb->', MO - MO.swapaxes(1, 3),
                                                                             MO, denom)
        return mp2_e