# This code is part of Qiskit.
#
# (C) Copyright IBM 2020, 2022.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE.txt file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""FCIDump parser."""

from typing import Any, Dict, Set, Tuple, Iterator
import itertools
import re
import numpy as np

from qiskit_nature import QiskitNatureError


def parse(fcidump: str) -> Dict[str, Any]:
    """Parses a FCIDump output.

    Args:
        fcidump: Path to the FCIDump file.
    Raises:
        QiskitNatureError: If the input file cannot be found, if a required field in the FCIDump
            file is missing, if wrong integral indices are encountered, or if the alpha/beta or
            beta/alpha 2-electron integrals are mixed.
    Returns:
        A dictionary storing the parsed data.
    """
    try:
        with open(fcidump, "r", encoding="utf8") as file:
            fcidump_str = file.read()
    except OSError as ex:
        raise QiskitNatureError(f"Input file '{fcidump}' cannot be read!") from ex

    output = {}  # type: Dict[str, Any]

    # FCIDump starts with a Fortran namelist of meta data
    namelist_end = re.search("(/|&END)", fcidump_str)
    metadata = fcidump_str[: namelist_end.start(0)]
    metadata = " ".join(metadata.split())  # replace duplicate whitespace and newlines
    # we know what elements to look for so we don't get too fancy with the parsing
    # pattern explanation:
    #  .*?      any text
    #  (*|*),   match either part of this group followed by a comma
    #  [-+]?    match up to a single - or +
    #  \d*.\d+  number format
    pattern = r".*?([-+]?\d*\.\d+|[-+]?\d+),"
    # we parse the values in the order in which they are listed in Knowles1989
    _norb = re.search("NORB" + pattern, metadata)
    if _norb is None:
        raise QiskitNatureError("The required NORB entry of the FCIDump format is missing!")
    norb = int(_norb.groups()[0])
    output["NORB"] = norb
    _nelec = re.search("NELEC" + pattern, metadata)
    if _nelec is None:
        raise QiskitNatureError("The required NELEC entry of the FCIDump format is missing!")
    output["NELEC"] = int(_nelec.groups()[0])
    # the rest of these values may occur and are set to their defaults otherwise
    _ms2 = re.search("MS2" + pattern, metadata)
    output["MS2"] = int(_ms2.groups()[0]) if _ms2 else 0
    _isym = re.search("ISYM" + pattern, metadata)
    output["ISYM"] = int(_isym.groups()[0]) if _isym else 1
    # ORBSYM holds a list, thus it requires a little different treatment
    _orbsym = re.search(r"ORBSYM.*?" + r"(\d+)," * norb, metadata)
    output["ORBSYM"] = [int(s) for s in _orbsym.groups()] if _orbsym else [1] * norb
    _iprtim = re.search("IPRTIM" + pattern, metadata)
    output["IPRTIM"] = int(_iprtim.groups()[0]) if _iprtim else -1
    _int = re.search("INT" + pattern, metadata)
    output["INT"] = int(_int.groups()[0]) if _int else 5
    _memory = re.search("MEMORY" + pattern, metadata)
    output["MEMORY"] = int(_memory.groups()[0]) if _memory else 10000
    _core = re.search("CORE" + pattern, metadata)
    output["CORE"] = float(_core.groups()[0]) if _core else 0.0
    _maxit = re.search("MAXIT" + pattern, metadata)
    output["MAXIT"] = int(_maxit.groups()[0]) if _maxit else 25
    _thr = re.search("THR" + pattern, metadata)
    output["THR"] = float(_thr.groups()[0]) if _thr else 1e-5
    _thrres = re.search("THRRES" + pattern, metadata)
    output["THRRES"] = float(_thrres.groups()[0]) if _thrres else 0.1
    _nroot = re.search("NROOT" + pattern, metadata)
    output["NROOT"] = int(_nroot.groups()[0]) if _nroot else 1

    # If the FCIDump file resulted from an unrestricted spin calculation the indices will label spin
    # rather than molecular orbitals. This means, that a line must exist which encodes the
    # coefficient for the spin orbital with index (norb*2, norb*2). By checking for such a line we
    # can distinguish between unrestricted and restricted FCIDump files.
    _uhf = bool(
        re.search(
            fr".*(\s+{norb * 2}\s+{norb * 2}\s+0\s+0)",
            fcidump_str[namelist_end.start(0) :],
        )
    )

    # the rest of the FCIDump will hold lines of the form x i a j b
    # a few cases have to be treated differently:
    # i, a, j and b are all zero: x is the core energy
    # TODO: a, j and b are all zero: x is the energy of the i-th MO  (often not supported)
    # j and b are both zero: x is the 1e-integral between i and a (x = <i|h|a>)
    # otherwise: x is the Coulomb integral ( x = (ia|jb) )
    hij = np.zeros((norb, norb))
    hij_elements = set(itertools.product(range(norb), repeat=2))
    hijkl = np.zeros((norb, norb, norb, norb))
    hijkl_elements = set(itertools.product(range(norb), repeat=4))
    hij_b = hijkl_ab = hijkl_ba = hijkl_bb = None
    hij_b_elements = hijkl_ab_elements = hijkl_ba_elements = hijkl_bb_elements = set()
    if _uhf:
        beta_range = [n + norb for n in range(norb)]
        hij_b = np.zeros((norb, norb))
        hij_b_elements = set(itertools.product(beta_range, repeat=2))
        hijkl_ab = np.zeros((norb, norb, norb, norb))
        hijkl_ba = np.zeros((norb, norb, norb, norb))
        hijkl_bb = np.zeros((norb, norb, norb, norb))
        hijkl_ab_elements = set(itertools.product(range(norb), range(norb), beta_range, beta_range))
        hijkl_ba_elements = set(itertools.product(beta_range, beta_range, range(norb), range(norb)))
        hijkl_bb_elements = set(itertools.product(beta_range, repeat=4))
    orbital_data = fcidump_str[namelist_end.end(0) :].split("\n")
    for orbital in orbital_data:
        if not orbital:
            continue
        x = float(orbital.split()[0])
        # Note: differing naming than ijkl due to E741 and this iajb is inline with this:
        # https://hande.readthedocs.io/en/latest/manual/integrals.html#fcidump-format
        i, a, j, b = [int(i) for i in orbital.split()[1:]]
        if i == a == j == b == 0:
            output["ecore"] = x
        elif a == j == b == 0:
            # TODO: x is the energy of the i-th MO
            continue
        elif j == b == 0:
            try:
                hij_elements.remove((i - 1, a - 1))
                hij[i - 1][a - 1] = x
            except KeyError as ex:
                if _uhf:
                    hij_b_elements.remove((i - 1, a - 1))
                    hij_b[i - 1 - norb][a - 1 - norb] = x
                else:
                    raise QiskitNatureError(
                        "Unkown 1-electron integral indices encountered in " f"'{(i, a)}'"
                    ) from ex
        else:
            try:
                hijkl_elements.remove((i - 1, a - 1, j - 1, b - 1))
                hijkl[i - 1][a - 1][j - 1][b - 1] = x
            except KeyError as ex:
                if _uhf:
                    try:
                        hijkl_ab_elements.remove((i - 1, a - 1, j - 1, b - 1))
                        hijkl_ab[i - 1][a - 1][j - 1 - norb][b - 1 - norb] = x
                    except KeyError:
                        try:
                            hijkl_ba_elements.remove((i - 1, a - 1, j - 1, b - 1))
                            hijkl_ba[i - 1 - norb][a - 1 - norb][j - 1][b - 1] = x
                        except KeyError:
                            hijkl_bb_elements.remove((i - 1, a - 1, j - 1, b - 1))
                            hijkl_bb[i - 1 - norb][a - 1 - norb][j - 1 - norb][b - 1 - norb] = x
                else:
                    raise QiskitNatureError(
                        "Unkown 2-electron integral indices encountered in " f"'{(i, a, j, b)}'"
                    ) from ex

    # iterate over still empty elements in 1-electron matrix and populate with symmetric ones
    # if any elements are not populated these will be zero
    _permute_1e_ints(hij, hij_elements, norb)

    if _uhf:
        # do the same for beta spin
        _permute_1e_ints(hij_b, hij_b_elements, norb, beta=True)

    # do the same of the 2-electron 4D matrix
    _permute_2e_ints(hijkl, hijkl_elements, norb)

    if _uhf:
        # do the same for beta spin
        _permute_2e_ints(hijkl_bb, hijkl_bb_elements, norb, beta=2)
        _permute_2e_ints(hijkl_ab, hijkl_ab_elements, norb, beta=1)  # type: ignore
        _permute_2e_ints(hijkl_ba, hijkl_ba_elements, norb, beta=1)  # type: ignore

        # assert that EITHER hijkl_ab OR hijkl_ba were given
        if np.allclose(hijkl_ab, 0.0) == np.allclose(hijkl_ba, 0.0):
            raise QiskitNatureError(
                "Encountered mixed sets of indices for the 2-electron \
                    integrals. Either alpha/beta or beta/alpha matrix should be specified."
            )

        if np.allclose(hijkl_ba, 0.0):
            hijkl_ba = hijkl_ab.transpose()

    output["hij"] = hij
    output["hij_b"] = hij_b
    output["hijkl"] = hijkl
    output["hijkl_ba"] = hijkl_ba
    output["hijkl_bb"] = hijkl_bb

    return output


def _permute_1e_ints(
    hij: np.ndarray, elements: Set[Tuple[int, ...]], norb: int, beta: bool = False
) -> None:
    for elem in elements.copy():
        shifted = tuple(e - (beta * norb) for e in elem)
        hij[shifted] = hij[shifted[::-1]]
        elements.remove(elem)


def _permute_2e_ints(
    hijkl: np.ndarray, elements: Set[Tuple[int, ...]], norb: int, beta: int = 0
) -> None:
    for elem in elements.copy():
        shifted = tuple(e - ((e >= norb) * norb) for e in elem)
        # initially look for "transposed" element if spins are equal
        if beta != 1 and elem[::-1] not in elements:
            hijkl[shifted] = hijkl[shifted[::-1]]
            elements.remove(elem)
            continue
        # then look at permutations of indices within the bra and ket respectively
        bra_perms = set(itertools.permutations(elem[:2]))
        ket_perms = set(itertools.permutations(elem[2:]))
        permutations: Iterator[Any]
        if beta == 1:
            # generally (ij|ab) != (ab|ij)
            # thus, the possible permutations are much less when the spins differ
            permutations = itertools.product(bra_perms, ket_perms)
        else:
            # ( ij | kl ) gives { ( ij | kl ), ( ij | lk ), ( ji | kl ), ( ji | lk ) }
            # AND { ( kl | ij ), ( kl | ji ), ( lk | ij ), ( lk | ji ) }
            # BUT NOT ( ik | jl ) etc.
            permutations = itertools.chain(
                itertools.product(bra_perms, ket_perms),
                itertools.product(ket_perms, bra_perms),
            )
        for perm in {e1 + e2 for e1, e2 in permutations}:
            if perm in elements:
                continue
            hijkl[shifted] = hijkl[tuple(e - (e >= norb) * norb for e in perm)]
            elements.remove(elem)
            break
