from pyscf.ao2mo import restore
from pyscf.tools import fcidump
import numpy as np
from numba import jit
import scipy


def scientific_superscript(num, digits=2):
    if np.isnan(num):
        print("WARNING: formatting of np.nan results in empty string")
        return ""
    base, exponent = f"{num:.{digits}e}".split("e")
    d = dict(zip("-+0123456789", "⁻⁺⁰¹²³⁴⁵⁶⁷⁸⁹"))
    return f'{base}⋅10{"".join(d.get(x, x) for x in  exponent)}'


def fcidump_to_integrals(fcidumpfile: str) -> tuple[float, np.ndarray, np.ndarray]:
    """
    Computes the 4-index Coulomb tensor from a FCIDUMP file

    :param fcidumpfile: string, path to FCIDUMP file
    :return: np.ndarray with 4 dimensions of equal length
    """
    data = fcidump.read(fcidumpfile)
    h0 = data["ECORE"]
    h1e = data["H1"]
    h2e = data["H2"]
    norb = data["NORB"]
    h2e = restore("s1", h2e, norb)
    return h0, h1e, h2e


def h1e_factorized(h1e: np.ndarray, h2e: np.ndarray) -> np.ndarray:
    """
    Yields the modified 1-electron coefficients of the hamiltonian when reordering
    the 2-electron term from

        H = Σ_pq,σ h_pq a_pσ^+ a_qσ  +  1/2 Σ_pqrs,στ h_pqrs a_pσ^+ a_rτ^+ a_sτ a_qσ
    to
        H = Σ_pq,σ h'_pq a_pσ^+ a_qσ  +  1/2 Σ_pqrs,στ h_pqrs a_pσ^+ a_qσ a_rτ^+ a_sτ .
    """
    h1e_fact = h1e - 1 / 2 * np.einsum("prrq->pq", h2e)
    return h1e_fact


def cholesky_pivoted(A: np.ndarray, tol=1e-12) -> np.ndarray:
    """
    A wrapper around lapack.dpstrf
    https://stackoverflow.com/a/77821449
    """
    n = A.shape[0]
    L, piv, rank, info = scipy.linalg.lapack.dpstrf(A, lower=1, tol=tol)

    # remove the garbage entries Lapack did not care about
    L = np.tril(L)
    L = L[:, :rank]

    # create the permutation matrix
    # maybe there's a more pythonic way to permute L?
    P = np.zeros((n, n))
    P[piv - 1, np.arange(n)] = 1
    L = P @ L
    return L


def h2e_cholesky(h2e: np.ndarray) -> np.ndarray:
    """
    Decomposes the 'matricised' Coulomb tensor into Cholesky vectors:

        h_pqrs = Σ_j L(j)_pq L(j)_rs

    Returns R matrices L(j) each of shape N x N, where R is the rank of the
    Coulomb tensor.
    """
    norb = h2e.shape[0]
    h2e = np.reshape(h2e, (norb**2, norb**2))
    cholesky_vecs = cholesky_pivoted(h2e)
    rank = cholesky_vecs.shape[1]
    h2e_factors = np.zeros((rank, norb, norb))
    for i in range(rank):
        h2e_factors[i, :, :] = np.reshape(cholesky_vecs[:, i], (norb, norb))
    return h2e_factors


def cholesky_vectors_to_h2e(cholesky_vecs: np.ndarray) -> np.ndarray:
    h2e = np.einsum("jpq,jrs->pqrs", cholesky_vecs, cholesky_vecs, optimize="optimal")
    return h2e


def factorized_to_standard_integrals(
    h1e_fact: np.ndarray, cholesky_vecs: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
    h1e = h1e_fact + 1 / 2 * np.einsum(
        "jpr,jrq", cholesky_vecs, cholesky_vecs, optimize="optimal"
    )
    h2e = cholesky_vectors_to_h2e(cholesky_vecs)
    return h1e, h2e


def h1e_1norm_cancellation(hamil):
    """
    Given a Hamiltonian H, define a 1-electron Hamiltonian H1 such that
    the 1-norm contribution of the quadratic terms in the remainder H - H1
    cancel exactly.

    :arg hamil: Hamiltonian H
    :return h1e: The coefficients of H1
    :rtype: np.ndarray
    """
    h1e = hamil[1]
    h2e = hamil[2]
    h1e_diag = h1e - 1 / 2 * np.einsum("prrq->pq", h2e, optimize="optimal")
    h1e_diag = h1e_diag + np.einsum("pqrr->pq", h2e, optimize="optimal")
    return h1e_diag


def opnorm_quadratic_hamil(h1e: np.ndarray, nelec: int | None = None) -> float:
    """
    Computes the operator norm (spectral norm) of a quadratic Hamiltonian in the nelec-particle
    subspace.
    Hamiltonian defined as
                            H = Σ_pq,σ h1e_pq a_pσ^+ a_qσ,
    σ is a spin-variable.
    """
    eigvals = np.linalg.eigvalsh(h1e)
    eigvals = sorted(2 * list(eigvals))[::-1]
    if nelec is None:
        sum_larger0 = max(np.cumsum(eigvals))
        sum_smaller0 = max(np.cumsum(-1 * np.array(eigvals)[::-1]))
    else:
        sum_larger0 = np.cumsum(eigvals)[nelec]
        sum_smaller0 = np.cumsum(-1 * np.array(eigvals)[::-1])[nelec]
    return max(sum_larger0, sum_smaller0)


def opnorm_cholesky_vec(L: np.ndarray, nelec: int | None = None) -> float:
    """
    Computes the operator norm of the hamiltonian term corresponding to a
    Cholesky-vector L(j) from the factorization of the Coulomb-tensor.
    Hamiltonian defined as
                            H = 1/2 Σ_pqrs,στ L(j)_pq L(j)_rs a_pσ^+ a_qσ a_rτ^+ a_sτ
    where σ and τ are spin variables.

    :param L: real-symmetric N x N matrix
    :param nelec: particle number, 0 <= nelec <= N

    :return: operator norm
    """
    return 1 / 2 * opnorm_quadratic_hamil(L, nelec) ** 2


@jit(nopython=True)
def norm1_paulisum(hamil) -> float:
    """
    Computes the 1-norm of the Hamiltonian after Fermion-to-qubit mapping, i.e.
    for resulting qubit Hamiltonian Σ_i h_i P_i, where P_i are Paulistrings, it returns
    λ = Σ_i |h_i|

    See equation F9 in https://quantum-journal.org/papers/q-2023-05-12-1000/
    """
    h1e = hamil[1]
    h2e = hamil[2]
    norm1 = 0
    norb = h1e.shape[0]

    for p in range(norb):
        for q in range(norb):
            curr = h1e[p, q]
            for r in range(norb):
                curr += h2e[p, q, r, r]
                curr -= 0.5 * h2e[p, r, r, q]
            norm1 += abs(curr)

    for r in range(norb - 1):
        for p in range(r + 1, norb):
            for q in range(norb - 1):
                for s in range(q + 1, norb):
                    norm1 += abs(h2e[p, q, r, s] - h2e[p, s, r, q]) / 2

    for r in range(norb - 1):
        for p in range(r + 1, norb):
            for q in range(norb):
                for s in range(q, norb):
                    norm1 += abs(h2e[p, q, r, s] / 2)

    for r in range(norb):
        for p in range(r, norb):
            for q in range(1, norb):
                for s in range(q):
                    norm1 += abs(h2e[p, q, r, s] / 2)

    for p in range(norb):
        for q in range(norb):
            norm1 += abs(h2e[p, q, p, q] / 4)

    return norm1


@jit(nopython=True)
def paulisum_hamil_coeffs(
    h0: float, h1e: np.ndarray, h2e: np.ndarray, include_constant=False
):
    """
    Computes the coefficients of the Hamiltonian after Fermion-to-qubit mapping, i.e.
    for resulting qubit Hamiltonian Σ_i h_i P_i, where P_i are Paulistrings, it returns
    the coefficients h_i.

    Returns tuple (coeff_constant, coeffs quadr terms, coeffs quart terms)

    See equation F9 in https://quantum-journal.org/papers/q-2023-05-12-1000/
    """
    coeff_const = 0
    coeffs_quadr = []
    coeffs_quart = []
    norb = h1e.shape[0]

    if include_constant:
        coeff_const = h0
        for p in range(norb):
            coeff_const += h1e[p, p]
            for r in range(norb):
                coeff_const += 0.5 * h2e[p, p, r, r]
                coeff_const -= 0.25 * h2e[p, r, r, p]

    for p in range(norb):
        for q in range(norb):
            curr = h1e[p, q]
            for r in range(norb):
                curr += h2e[p, q, r, r]
                curr -= 0.5 * h2e[p, r, r, q]
            coeffs_quadr.append(1j / 2 * curr)
            coeffs_quadr.append(1j / 2 * curr)

    for r in range(norb - 1):
        for p in range(r + 1, norb):
            for q in range(norb - 1):
                for s in range(q + 1, norb):
                    coeffs_quart.append((h2e[p, q, r, s] - h2e[p, s, r, q]) / 4)
                    coeffs_quart.append((h2e[p, q, r, s] - h2e[p, s, r, q]) / 4)

    for r in range(norb - 1):
        for p in range(r + 1, norb):
            for q in range(norb):
                for s in range(q, norb):
                    coeffs_quart.append(h2e[p, q, r, s] / 4)
                    coeffs_quart.append(h2e[p, q, r, s] / 4)

    for r in range(norb):
        for p in range(r, norb):
            for q in range(1, norb):
                for s in range(q):
                    coeffs_quart.append(h2e[p, q, r, s] / 4)
                    coeffs_quart.append(h2e[p, q, r, s] / 4)

    for p in range(norb):
        for q in range(norb):
            coeffs_quart.append(h2e[p, q, p, q] / 4)

    return coeff_const, np.array(coeffs_quadr), np.array(coeffs_quart)