# Copyright (c) 2024, the cclib development team
#
# This file is part of cclib (http://cclib.github.io) and is distributed under
# the terms of the BSD 3-Clause License.

"""Parser for Formatted Checkpoint files"""


from cclib.parser import logfileparser, utils

import numpy

SHELL_ORBITALS = {
    0: ["S"],
    1: ["PX", "PY", "PZ"],
    -1: ["S", "PX", "PY", "PZ"],
    2: ["D1", "D2", "D3", "D4", "D5", "D6"],
    -2: ["D1", "D2", "D3", "D4", "D5"],
    3: ["F1", "F2", "F3", "F4", "F5", "F6", "F7", "F8", "F9", "F10"],
    -3: ["F1", "F2", "F3", "F4", "F5", "F6", "F7"],
    4: ["G1", "G2", "G3", "G4", "G5", "G6", "G7", "G8", "G9", "G10", "G11", "G12", "G13"],
    -4: ["G1", "G2", "G3", "G4", "G5", "G6", "G7", "G8", "G9"],
}

SHELL_START = {0: 1, 1: 2, -1: 2, 2: 3, -2: 3, 3: 4, -3: 4}


def _shell_to_orbitals(type, offset):
    """Convert a Fchk shell type and offset to a list of string representations.

    For example, shell type = -2 corresponds to d orbitals (spherical basis) with
    an offset = 1 would correspond to the 4d orbitals, so this function returns
    `['4D1', '4D2', '4D3', '4D4', '4D5']`.
    """

    return [f"{SHELL_START[type] + offset}{x}" for x in SHELL_ORBITALS[type]]


class FChk(logfileparser.Logfile):
    """A Formatted checkpoint file, which contains molecular and wavefunction information.

    These files are produced by Gaussian and Q-Chem.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(logname="FChk", *args, **kwargs)
        self.start = True

    def __str__(self):
        """Return a string representation of the object."""
        return f"Formatted checkpoint file {self.filename}"

    def __repr__(self):
        """Return a representation of the object."""
        return f'FCHK("{self.filename}")'

    def normalisesym(self, symlabel):
        """Just return label"""
        return symlabel

    def extract(self, inputfile, line):
        # just opened file, skip first line to get basis
        if self.start:
            method = next(inputfile)
            self.metadata["basis_set"] = method.split()[-1]
            self.start = False

        if line[0:6] == "Charge":
            self.set_attribute("charge", int(line.split()[-1]))

        if line[0:12] == "Multiplicity":
            self.set_attribute("mult", int(line.split()[-1]))

        if line[0:14] == "Atomic numbers":
            self.natom = int(line.split()[-1])
            atomnos = self._parse_block(inputfile, self.natom, int, "Basic Information")
            self.set_attribute("atomnos", atomnos)

        if line[0:19] == "Number of electrons":
            alpha = next(inputfile)
            alpha_homo = int(alpha.split()[-1]) - 1

            beta = next(inputfile)
            beta_homo = int(beta.split()[-1]) - 1

            self.set_attribute("homos", [alpha_homo, beta_homo])

        if line[0:29] == "Current cartesian coordinates":
            count = int(line.split()[-1])
            assert count % 3 == 0

            coords = numpy.array(self._parse_block(inputfile, count, float, "Coordinates"))
            coords.shape = (1, int(count / 3), 3)
            self.set_attribute("atomcoords", utils.convertor(coords, "bohr", "Angstrom"))

        if line[0:10] == "SCF Energy":
            self.set_attribute(
                "scfenergies", [utils.convertor(float(line.split()[-1]), "hartree", "eV")]
            )

        if line[0:25] == "Number of basis functions":
            self.set_attribute("nbasis", int(line.split()[-1]))

        if line[0:14] == "Overlap Matrix":
            count = int(line.split()[-1])
            self.set_attribute(
                "aooverlaps",
                utils.block_to_matrix(
                    numpy.asarray(self._parse_block(inputfile, count, float, "Overlap Matrix"))
                ),
            )

        if line[0:31] == "Number of independent functions":
            self.set_attribute("nmo", int(line.split()[-1]))

        if line[0:21] == "Alpha MO coefficients":
            count = int(line.split()[-1])
            assert count == self.nbasis * self.nmo

            coeffs = numpy.array(self._parse_block(inputfile, count, float, "Alpha Coefficients"))
            coeffs.shape = (self.nmo, self.nbasis)
            self.set_attribute("mocoeffs", [coeffs])

        if line[0:22] == "Alpha Orbital Energies":
            count = int(line.split()[-1])
            assert count == self.nmo

            energies = numpy.array(self._parse_block(inputfile, count, float, "Alpha MO Energies"))
            self.set_attribute("moenergies", [utils.convertor(energies, "hartree", "eV")])

        if line[0:20] == "Beta MO coefficients":
            count = int(line.split()[-1])
            assert count == self.nbasis * self.nmo

            coeffs = numpy.array(self._parse_block(inputfile, count, float, "Beta Coefficients"))
            coeffs.shape = (self.nmo, self.nbasis)
            self.append_attribute("mocoeffs", coeffs)

        if line[0:21] == "Beta Orbital Energies":
            count = int(line.split()[-1])
            assert count == self.nmo

            energies = numpy.array(self._parse_block(inputfile, count, float, "Alpha MO Energies"))
            self.append_attribute("moenergies", utils.convertor(energies, "hartree", "eV"))

        if line[0:11] == "Shell types":
            self.parse_aonames(line, inputfile)

        if line[0:19] == "Real atomic weights":
            count = int(line.split()[-1])
            assert count == self.natom

            atommasses = numpy.array(self._parse_block(inputfile, count, float, "Atomic Masses"))

            self.set_attribute("atommasses", atommasses)

        if line[0:10] == "SCF Energy":
            self.scfenergy = float(line.split()[-1])

            self.set_attribute("scfenergies", [utils.convertor(self.scfenergy, "hartree", "eV")])

        if line[0:16] == "Mulliken Charges":
            count = int(line.split()[-1])
            if not hasattr(self, "atomcharges"):
                self.atomcharges = {}
            self.atomcharges["mulliken"] = self._parse_block(
                inputfile, count, float, "Mulliken Charges"
            )

        if line[0:18] == "Cartesian Gradient":
            count = int(line.split()[-1])
            assert count == self.natom * 3

            gradient = numpy.array(self._parse_block(inputfile, count, float, "Gradient"))

            self.set_attribute("grads", gradient)

        if line[0:14] == "Polarizability":
            polarizability = numpy.asarray(self._parse_block(inputfile, 6, float, "Polarizability"))
            self.append_attribute("polarizabilities", utils.block_to_matrix(polarizability))

        if line[0:25] == "Cartesian Force Constants":
            count = int(line.split()[-1])
            assert count == (3 * self.natom * (3 * self.natom + 1)) / 2

            hessian = numpy.array(self._parse_block(inputfile, count, float, "Hessian"))

            self.set_attribute("hessian", utils.block_to_matrix(hessian))

        if line[0:22] == "Number of Normal Modes":
            self.set_attribute("nmode", int(line.split()[-1]))

        if line[0:6] == "Vib-E2":
            count = int(line.split()[-1])
            le2 = self._parse_block(inputfile, count, float, "Vib-E2")
            assert hasattr(self, "nmode")
            nmode = self.nmode
            self.set_attribute("vibfreqs", le2[:nmode])
            le2 = le2[nmode:]
            self.set_attribute("vibrmasses", le2[:nmode])
            le2 = le2[nmode:]
            self.set_attribute("vibfconsts", le2[:nmode])
            le2 = le2[nmode:]
            self.set_attribute("vibirs", le2[:nmode])
            # The rest should be empty unless Raman or ROA were calculated.
            le2 = numpy.asarray(le2[nmode:])
            if numpy.any(le2[:nmode] >= 1.0e-20):
                self.set_attribute("vibramans", le2[:nmode])
                le2 = le2[nmode:]

        if line[0:9] == "Vib-Modes":
            count = int(line.split()[-1])
            vibdisps = numpy.asarray(self._parse_block(inputfile, count, float, "Vib-Modes"))
            # indices from fast to slow are [xyz, atom, mode]
            assert hasattr(self, "nmode")
            assert count == 3 * self.natom * self.nmode
            self.set_attribute("vibdisps", vibdisps.reshape(self.nmode, self.natom, 3))

        if line[0:13] == "ETran scalars":
            count = int(line.split()[-1])

            etscalars = self._parse_block(inputfile, count, int, "ET Scalars")

            # Set attribute: self.netroot (number of excited estates)
            self.netroot = etscalars[4]

        if line[0:10] == "ETran spin":
            count = int(line.split()[-1])

            etspin = self._parse_block(inputfile, count, int, "ET Spin")

            spin_labels = {0: "Singlet", 2: "Triplet", -1: "Unknown"}
            etsyms = []
            for i in etspin:
                if i in spin_labels:
                    etsyms.append(spin_labels[i])
                else:
                    etsyms.append(spin_labels[-1])

            # The extracted property does not contain the actual irrep label
            # (contrarily to that extracted from the Gaussian log)
            # After this, 'Etran sym' appears (and would need to be parsed),
            # but at least in Gaussian this contains only zeroes regardless of the irrep.

            self.set_attribute("etsyms", etsyms)

        if line[0:18] == "ETran state values":
            # This section is organized as follows:
            # ·First the properties of each excited state (up to net):
            # E, {muNx,muNy,muNz,muvelNx,muvelNy,muvelNz,mmagNx,mmagNy,mmagNz,unkX,unkY,unkZ,unkX,unkY,unkZ}_N=1,net
            # ·Then come 48 items (only if Freq is requested)
            # They were all 0.000 in G09, but get an actual value in G16
            # ·Then, the derivates of each property with respect to Cartesian coordiates only for target state (netroot)
            # For each Cartesian coordiate, all derivatives wrt to it are included:
            #  dE/dx1 dmux/dx1 dmuy/dx1 ... unkZ/dx1
            #  dE/dy1 dmux/dy1 dmuy/dy1 ... unkZ/dy1
            #  ...
            #  dE/dzN dmux/dzN dmuy/dzN ... unkZ/dzN
            # The number of items is therefore:
            ### 16*net (no Freq jobs)
            ### 16*net + 48 + 3*self.natom*16 (Freq jobs)
            count = int(line.split()[-1])
            if hasattr(self, "etsyms"):
                net = len(self.etsyms)
            else:
                net = 0  # This forces an AssertionError below
            assert count in [16 * net, 16 * net + 48 + 3 * self.natom * 16]

            etvalues = self._parse_block(inputfile, count, float, "ET Values")

            # ETr energies (1/cm)
            etenergies_au = [e_es - self.scfenergy for e_es in etvalues[0 : net * 16 : 16]]
            etenergies = [utils.convertor(etr, "hartree", "wavenumber") for etr in etenergies_au]
            self.set_attribute("etenergies", etenergies)

            # ETr dipoles (length-gauge)
            etdips = []
            for k in range(1, 16 * net, 16):
                etdips.append(etvalues[k : k + 3])
            self.set_attribute("etdips", etdips)

            # Osc. Strength from Etr dipoles
            # oscs = 2/3 * Etr(au) * dip²(au)
            etoscs = [
                2.0 / 3.0 * e * numpy.linalg.norm(numpy.array(dip)) ** 2
                for e, dip in zip(etenergies_au, etdips)
            ]
            self.set_attribute("etoscs", etoscs)

            # ETr dipoles (velocity-gauge)
            etveldips = []
            for k in range(4, 16 * net, 16):
                etveldips.append(etvalues[k : k + 3])
            self.set_attribute("etveldips", etveldips)

            # ETr magnetic dipoles
            etmagdips = []
            for k in range(7, 16 * net, 16):
                etmagdips.append(etvalues[k : k + 3])
            self.set_attribute("etmagdips", etmagdips)

        if line[0:19] == "Excitation Energies":
            count = int(line.split()[-1])
            etenergies = self._parse_block(inputfile, count, float, "Excitation energies")
            etenergies = [utils.convertor(e, "hartree", "wavenumber") for e in etenergies]
            self.set_attribute("etenergies", etenergies)

        if line[0:20] == "Oscillator Strengths":
            count = int(line.split()[-1])
            etoscs = self._parse_block(inputfile, count, float, "Oscillator Strengths")
            self.set_attribute("etoscs", etoscs)

    def parse_aonames(self, line, inputfile):
        # e.g.: Shell types                                I   N=          28
        count = int(line.split()[-1])
        shell_types = self._parse_block(inputfile, count, int, "Atomic Orbital Names")

        # e.g.: Number of primitives per shell             I   N=          28
        next(inputfile)
        self._parse_block(inputfile, count, int, "Atomic Orbital Names")

        # e.g. Shell to atom map                          I   N=          28
        next(inputfile)
        shell_map = self._parse_block(inputfile, count, int, "Atomic Orbital Names")

        elements = (self.table.element[x] for x in self.atomnos)
        atom_labels = [f"{y}{x}" for x, y in enumerate(elements, 1)]

        # get orbitals for first atom and start aonames and atombasis lists
        atom = shell_map[0] - 1
        shell_offset = 0
        orbitals = _shell_to_orbitals(shell_types[0], shell_offset)
        aonames = [f"{atom_labels[atom]}_{x}" for x in orbitals]
        atombasis = [list(range(len(orbitals)))]

        # get rest
        for i in range(1, len(shell_types)):
            _type = shell_types[i]
            atom = shell_map[i] - 1
            shell_offset += 1
            basis_offset = (
                atombasis[-1][-1] + 1
            )  # atombasis is increasing numbers, so just grab last

            # if we've move to next atom, need to update offset of shells (e.g. start at 1S)
            # and start new list for atom basis
            if atom != shell_map[i - 1] - 1:
                shell_offset = 0
                atombasis.append([])

            # determine if we've changed shell type (e.g. from S to P)
            if _type != shell_types[i - 1]:
                shell_offset = 0

            orbitals = _shell_to_orbitals(_type, shell_offset)
            aonames.extend([f"{atom_labels[atom]}_{x}" for x in orbitals])
            atombasis[-1].extend(list(range(basis_offset, basis_offset + len(orbitals))))

        assert (
            len(aonames) == self.nbasis
        ), f"Length of aonames != nbasis: {len(aonames)} != {self.nbasis}"
        self.set_attribute("aonames", aonames)

        assert (
            len(atombasis) == self.natom
        ), f"Length of atombasis != natom: {len(atombasis)} != {self.natom}"
        self.set_attribute("atombasis", atombasis)

    def after_parsing(self):
        """Correct data or do parser-specific validation after parsing is finished."""

        # If restricted calculation, need to remove beta homo
        if len(self.moenergies) == len(self.homos) - 1:
            self.homos.pop()

    def _parse_block(self, inputfile, count, type, msg):
        atomnos = []
        while len(atomnos) < count:
            self.updateprogress(inputfile, msg, self.fupdate)
            line = next(inputfile)
            atomnos.extend([type(x) for x in line.split()])
        return atomnos
