from warnings import warn
import itertools
import logging
from scipy.special import erfc
import numpy as np
from numba import jit
from dopyqo.colors import *


def nuclear_repulsion_energy_ewald(
    atom_positions: np.ndarray,
    atomic_numbers: np.ndarray,
    lattice_vectors: np.ndarray,
    lattice_vectors_reciprocal: np.ndarray,
    cell_volume: float,
    gcut: float,
    sigma: float | None = None,
    rust_impl: bool = True,
) -> np.ndarray:
    r"""Calculate the nuclear repulsion energy E between all nuclei in a periodic lattice, i.e.,
    E = 1/2 \sum_I \sum_J \sum_T' Z_I Z_J / |R_I-R_J-T|
    where ' denotes that the term T=0 and I=J is omitted. Z_I are the charges of the nuclei
    and R_I their positions. T are all possible lattice translation vectors.
    This is equivalent to
    E = 1/2 \sum_I Z_I \Phi_I(R_I)
    where \Phi_I(r) is the electrostatic potential generated by all periodic nuclei EXCEPT
    the nucleus Z_I at position R_I (its periodic images are still included).
    To calculate \Phi_I(r) we define charge densities and solve the Poisson equation in
    real- and reciprocal-space. The nuclei Z_I are modelled by delta distribution at positions R_I+T.
    Additionally, Ewalds method is used, i.e., we add and substract Gaussian charge distributions
    at all R_I+T, and subtract a homogenous background charge simulating the electron-charges.
    A short-ranged term including all delta distributions minus all Gaussian charges (except the
    term where T=0 and I=J) can be calculated in real-space. A long-ranged term including
    all Gaussian charges at R_I+T minus the constant background charge can be calculated
    in reciprocal-space. The term describing the subtraction of the Gaussian charge at R_I
    is calculated in real-space and is the so-called self-energy term. The G -> 0 limit
    in the reciprocal space calculation is also calculated separately and is called the
    charged-system term. The value \sigma determining the splitting between the real-space
    and reciprocal-space term is calculated as done in Quantum ESPRESSO (see PW/src/ewald.f90:
    https://github.com/QEF/q-e/blob/de3035747f5d8f2ec9a67869827341ebb43f5b12/PW/src/ewald.f90)

    We calculate E = E_S + E_L + E_self + E_charged with

    Short-ranged term calculated in real-space
    E_S    = 1/2 \sum_T' \sum_I \sum_J Z_I Z_J / |R_I - R_J - T| erfc(|R_I - R_J - T| \sqrt(\sigma))
    where ' denotes that the term T=0 and I=J is omitted.

    Long-ranged term calculated in reciprocal-space
    E_L    = 1/2 4\pi/V \sum_{k \neq 0} \sum_I \sum_J Z_I Z_J / |k|^2 e^{i k . (R_I - R_J)} e^{- 1/(4 \sigma) |k|^2}

    Self-energy originated in subtracting the Gaussian charge at position R_I from the periodic charge density
    generated by all atoms (in real-space)
    E_self = - \sum_I (Z_I^2) \sqrt(\sigma) / \sqrt(\pi)

    Charged-system term from the G=0 term in the reciprocal-space sum
    E_charged = - \pi / (2 V \sigma)  (\sum_I Z_I)^2

    Notes:
    - When pseudopotentials are used Z_I are the valence charges, i.e. the charge of the nucleus that is not
      modelled with the pseudopotential.
    - This function should calculate the same Ewald energy calculated in Quantum ESPRESSO, called
      "ewald contribution" in the pw.x output file and "ewald" in the total_energy section in
      the output-xml.
    - Equivalent to Ewald calculation in PySCF of a cell in pyscf.pbc.get.cell.ewald function

    References:
    - http://metal.elte.hu/~groma/Anyagtudomany/kittel.pdf Appendix B
    - https://journals.aps.org/prb/pdf/10.1103/PhysRevB.53.1814, equation (16)
    - https://courses.physics.illinois.edu/phys466/sp2013/projects/2003/Team2/ewald_text.htm,
    - S. W. de Leeuw et al. Proc. R. Soc. Lond. A 1980 373, 27-56, doi: 10.1098/rspa.1980.0135,  equation (1.5)
    - M. P. Allen and D. J. Tildesley "Computer Simulation of Solids", equation (5.20)
    - http://micro.stanford.edu/mediawiki/images/4/46/Ewald_notes.pdf
    - https://homepages.uc.edu/~becktl/tlb_ewald.pdf
    - Quantum ESPRESSO source code in q-e/PW/src/ewald.f90

    Args:
        atom_positions (np.ndarray): Array of the coordinates R_I of every atom described
                                     by the pseudopotential. Shape (#atoms, 3)
        atomic_numbers (np.ndarray):  1D array of atomic number of each atom. Shape (#atoms,)
        lattice_vectors (np.ndarray): 2D array of lattice vectors. Each row is one lattice vector. Shape (3, 3)
        lattice_vectors_reciprocal (np.ndarray): 2D array of reciprocal lattice vectors. Each row is one reciprocal lattice vector. Shape (3, 3)
        cell_volume (float): Cell volume of the computationen real space cell.
        gcut (float): Cutoff of the reciprocal space sum in E_L, i.e., k < gcutrho. In Quantum ESPRESSO the value of the G-cutoff of the density is used.
        sigma (float | None, optional): Value for \sigma used in the Ewald summation to split real- and reciprocal-space sums. If None this is calculated to satisfy
                                        \sum_I 2 Z_I \sqrt{\frac\sigma\pi} \mathrm{erfc}\left(\sqrt{\frac{\left(4G_\mathrm{cut}\right)^2}{4\sigma}}\right)  \leq 10^{-7}
                                        Defaults to None.

    Raises:
        RuntimeError: If \sigma could not be calculated.

    Returns:
        float: Ewald energy
    """

    if rust_impl:
        try:
            import dopyqo_rs as calc_rs
        except ImportError:
            print(f"{ORANGE}Ewald warning: Could not import dopyqo_rs package. Falling back to python implementation.{RESET_COLOR}")
            rust_impl = False  # Set such that variable accurately represents if rust implementation is used
        else:  # No exception
            logging.info("Using Rust implementation.")
            return calc_rs.ewald(
                atom_positions,
                atomic_numbers,
                lattice_vectors,
                lattice_vectors_reciprocal,
                cell_volume,
                gcut,
            )

    if sigma is None:
        sigma = 2.8
        charge = np.sum(atomic_numbers)
        gcutm = gcut**2
        #
        # choose sigma in order to have convergence in the sum over G
        # upperbound is a safe upper bound for the error in the sum over G
        #
        while True:
            if sigma <= 0.0:
                raise RuntimeError("Optimal sigma for Ewald sum not found!")
            upperbound = 2.0 * charge**2 * np.sqrt(sigma / np.pi) * erfc(np.sqrt(gcutm / 4.0 / sigma))
            if upperbound > 1e-7:
                sigma = sigma - 0.1
            else:
                break

    logging.info("sigma %f", sigma)

    e_short = 0.0  # real-space sum
    converged_real = False
    # n_max = 5
    n_max_used = 0

    alat = np.linalg.norm(lattice_vectors[0], ord=2)
    t_vec_max_norm = 4.0 / np.sqrt(sigma) / alat
    b_norms = np.linalg.norm(lattice_vectors_reciprocal, ord=2, axis=1)
    n_max_x, n_max_y, n_max_z = (b_norms * t_vec_max_norm).astype(int) + 2

    # Generate list or translation vectors ordered by their norm
    # n_unordered = itertools.product(range(-n_max, n_max + 1), repeat=3)
    logging.debug("Generating unordered translation vectors...")
    t_vecs_unordered = [
        [np.dot(lattice_vectors, np.array([nx, ny, nz])), [nx, ny, nz]]
        for nx, ny, nz in itertools.product(
            range(-n_max_x, n_max_x + 1),
            range(-n_max_y, n_max_y + 1),
            range(-n_max_z, n_max_z + 1),
        )
        # for nx, ny, nz in n_unordered
    ]

    t_vecs = sorted(t_vecs_unordered, key=lambda x: np.linalg.norm(x[0], ord=2))

    for t_vec, (nx, ny, nz) in t_vecs:
        for i, pos_i in enumerate(atom_positions):  # \sum_I
            z_i = atomic_numbers[i]
            for j, pos_j in enumerate(atom_positions):  # \sum_J
                if np.all(t_vec == 0.0) and i == j:
                    continue
                z_j = atomic_numbers[j]
                r = np.linalg.norm(pos_i - pos_j - t_vec, ord=2)
                e_short += z_i * z_j * erfc(r * np.sqrt(sigma)) / r

                # with this choice terms up to ZiZj*erfc(4) are counted (erfc(4)=1.5e-8)
                if r >= 4 / np.sqrt(sigma):
                    # print(f"Stopping real-sum at {nx}/{ny}/{nz} (sigma={sigma})")
                    converged_real = True
                    n_max_used = np.max([np.abs(nx), np.abs(ny), np.abs(nz), n_max_used])
                    # break
        # if converged_real:
        #     break
    if converged_real:
        logging.info("Real-space sum converged! Largest n was %i", n_max_used)
    else:
        logging.warning("Real-space sum NOT converged!")

    e_short *= 1 / 2

    e_long = 0.0  # reciprocal-space sum

    # Estimate size of Miller indices grid using the cutoff-energy and reciprocal lattice vectors
    # A simple, but not accurate enough, estimate would be gcutrho/|b_i| for i \in {1, 2, 3}
    # Here we take the shape of the reciprocal lattice into account:
    #       max[ gcutrho/(|b_i| + \sum_{j \neq i} b_i . b_j ) ] for i \in {1, 2, 3}
    # TODO: To use the same grid QE uses we have to read the Miller indices used for the density,
    #       which is probably saved in charge-density.dat.
    #       Alternatively check how pymatgen chooses the reciprocal grid in its Ewald calculation method
    b1 = lattice_vectors_reciprocal[0]
    b2 = lattice_vectors_reciprocal[1]
    b3 = lattice_vectors_reciprocal[2]
    mill_max = max(
        [
            round(abs(gcut / (np.linalg.norm(b1) + np.dot(b1, b2) + np.dot(b1, b3)))),
            round(abs(gcut / (np.linalg.norm(b2) + np.dot(b2, b1) + np.dot(b2, b3)))),
            round(abs(gcut / (np.linalg.norm(b3) + np.dot(b3, b1) + np.dot(b3, b2)))),
        ]
    )
    x, y, z = np.meshgrid(
        np.arange(-mill_max, mill_max + 1),
        np.arange(-mill_max, mill_max + 1),
        np.arange(-mill_max, mill_max + 1),
        indexing="ij",
    )
    mill_rho = np.stack((x.ravel(), y.ravel(), z.ravel()), axis=1)
    k_vecs = np.einsum("ij, kj -> ki", lattice_vectors_reciprocal.T, mill_rho)
    norms = np.linalg.norm(k_vecs, axis=1)
    k_vecs = k_vecs[norms <= gcut]
    for k_vec in k_vecs:
        k2 = np.linalg.norm(k_vec, ord=2) ** 2
        if np.isclose(k2, 0.0):
            continue
        structure_factor = np.sum(atomic_numbers * np.exp(1.0j * np.dot(atom_positions, k_vec)))
        e_long += np.exp(-k2 / (4 * sigma)) * np.abs(structure_factor) ** 2 / k2

    e_long *= 4.0 * np.pi / 2.0 / cell_volume

    logging.debug("Reciprocal-space sum calculated!")

    # Term arising from substracting the Gaussian charge distribution at R_I
    # in real-space
    e_self = -np.sum(atomic_numbers**2) * np.sqrt(sigma / np.pi)

    # The analytic limit for G → 0 in the reciprocal sum
    # The G=0 term in the reciprocal term is only partly cancelled by the uniform
    # background charge. The remaining term is the following
    # Also see Martin, 2020, "Electronic Structure" eq. (F.6) at the text afterwards
    # which unfortunately misses a derivation
    e_charged = -np.pi / cell_volume / sigma * np.sum(atomic_numbers) ** 2 / 2
    # e_charged = -np.sum(atomic_numbers) ** 2 / sigma / 2 * np.pi / cell_volume

    e_tot = e_short + e_long + e_self + e_charged

    logging.info(
        "e_short: %f, e_long: %f, e_self: %f, e_charged: %f, e_tot: %f",
        e_short,
        e_long,
        e_self,
        e_charged,
        e_tot,
    )

    return e_tot


def check_symmetry_one_body_matrix(matrix: np.ndarray):
    """Check if given matrix satisfies the symmetries of one-body matrices (hermitian)

    Args:
        matrix (np.ndarray): Matrix to check symmetries of

    Returns:
        bool: Bool specifying of symmetry is satisfied
    """
    matrix_hermitian = matrix.T.conj()
    allclose_hermitian = np.allclose(matrix, matrix_hermitian)
    return allclose_hermitian


def check_symmetry_two_body_matrix(matrix: np.ndarray):
    """Check if given matrix satisfies the symmetries of ERIs

    Args:
        matrix (np.ndarray): Matrix to check symmetries of

    Returns:
        tuple[bool, bool, bool]: Bools specifying if following symmetries are satisfied: swap symmetry, hermitian symmetry, hermitian+swap symmetry
    """
    matrix_swap = matrix.swapaxes(0, 1).swapaxes(2, 3)
    matrix_hermitian = matrix.T.conj()
    matrix_hemitian_swap = matrix.T.conj().swapaxes(0, 1).swapaxes(2, 3)

    allclose_swap = np.allclose(matrix, matrix_swap)  # , rtol=1e-5, atol=1e-5)
    allclose_hermitian = np.allclose(matrix, matrix_hermitian)  # , rtol=1e-5, atol=1e-5)
    allclose_hermitian_swap = np.allclose(matrix, matrix_hemitian_swap)  # , rtol=1e-5, atol=1e-5)

    allclose_ljki = None
    allclose_ikjl = None
    allclose_kilj = None
    allclose_jlik = None
    if np.allclose(matrix.real, matrix):
        logging.info("Two-body matrix is real. Checking real-symmetries...")
        allclose_hermitian = np.allclose(matrix, matrix.T)  # ijkl = lkji
        allclose_swap = np.allclose(matrix, matrix.transpose(1, 0, 3, 2))  # ijkl = jilk
        allclose_ljki = np.allclose(matrix, matrix.transpose(3, 1, 2, 0))  # ijkl = ljki
        allclose_ikjl = np.allclose(matrix, matrix.transpose(0, 2, 1, 3))  # ijkl = ikjl
        allclose_hermitian_swap = np.allclose(matrix, matrix.T.transpose(1, 0, 3, 2))  # ijkl = klij
        allclose_kilj = np.allclose(matrix, matrix.transpose(2, 0, 3, 1))  # ijkl = kilj
        allclose_jlik = np.allclose(matrix, matrix.transpose(1, 3, 0, 2))  # ijkl = jlik

    return (
        allclose_swap,
        allclose_hermitian,
        allclose_hermitian_swap,
        allclose_ljki,
        allclose_ikjl,
        allclose_kilj,
        allclose_jlik,
    )


def iTj(p: np.ndarray, c_ip: np.ndarray) -> np.ndarray:  # Calculates <i|T|j>
    """Calculate kinetic energy matrix in Hartree units in the Kohn-Sham basis

    Args:
        p (np.ndarray): Array of momentum vectors, shape (#waves, 3)
        c_ip (np.ndarray): Array of coefficients describing the Kohn-Sham orbitals
            in the plane wave basis, shape (#bands, #waves)

    Returns:
        np.ndarray: Kinetic energy matrix
    """
    # Kinetic energy matrix in Hartree units
    p_norm = p[:, 0] ** 2 + p[:, 1] ** 2 + p[:, 2] ** 2  # p^2
    # 1/2 <i|p^2|j>
    return 0.5 * np.einsum("ip, p, jp -> ij", c_ip.conjugate(), p_norm, c_ip)


def iUj(
    p: np.ndarray,
    c_ip: np.ndarray,
    atom_positions: np.ndarray,
    atomic_numbers: np.ndarray,
    cell_volume: float,
) -> np.ndarray:
    r"""Calculate nuclear interaction (U(r) = \sum_i Z_i / |R_I-r|) matrix in Hartree units in the Kohn-Sham basis

    Args:
        p (np.ndarray): Array of momentum vectors
        c_ip (np.ndarray): Array of coefficients describing the Kohn-Sham orbitals in the plane wave basis
        atom_positions (np.ndarray): 2D array of positions of each atom. Shape (#atoms, 3)
        atomic_numbers (np.ndarray): 1D array of atomic number of each ato. Shape (#atoms,)m
        cell_volume (float): Cell volume of the computationen real space cell.

    Returns:
        np.ndarray: Nuclear interaction matrix
    """
    # Nuclear interaction matrix in Hartree units
    # Calculates <i|U|j>
    # <p|U|q> = 4\pi / cell_volume \sum_I Z_I  exp(-i (p-q) . R_I) 1/(q-p)²
    # <i|U|j> = \sum_{p,q,p!=q}  c_{p,i}^* c_{q,j} U_pq = 4\pi / cell_volume \sum_{p,q,p!=q} \sum_I Z_I  c_{p,i}^* c_{q,j} exp(-i (p-q) . R_I) 1/(q-p)²
    q = p
    Z_I = atomic_numbers  # shape (#atoms,)
    R_I = atom_positions  # shape (#atoms, 3)

    q = p
    p_minus_q = p[:, None] - q[None]  # shape (#waves, #waves, 3)
    # q[None] is of shape (1, #waves, 3) while p[:, None] is of shape (#waves, 1, 3)
    # so p_minus_q[i,j] = p[i] - q[j]

    p_minus_q = q[:, None] - p[None]  # shape (#waves, #waves, 3)
    p_minus_q_norm = np.linalg.norm(p_minus_q, ord=2, axis=2)  # Has zeros on diagonal, shape (#waves, #waves)
    p_minus_q_norm_filtered = p_minus_q_norm.copy() + np.eye(p_minus_q_norm.shape[0])  # Replace zeros on diagonal with ones

    p_minus_q_dot_R = np.sum(p_minus_q[None] * R_I[:, None, None], axis=3)  # sum over 3D-coordinates, shape (#atoms, #waves, #waves)
    # \sum_I 4\pi/V e^{-i (G-G').R_I}
    prefactor = 4 * np.pi / cell_volume * np.sum(Z_I[:, None, None] * np.exp(-1j * p_minus_q_dot_R), axis=0)  # sum over atoms, shape (#waves, #waves)
    one_over_p_minus_q_norm_squared = 1 / p_minus_q_norm_filtered**2
    one_over_p_minus_q_norm_squared = one_over_p_minus_q_norm_squared.copy() - np.eye(one_over_p_minus_q_norm_squared.shape[0])

    u_mat = prefactor * one_over_p_minus_q_norm_squared
    assert (u_mat.diagonal() == 0.0).all()

    return c_ip.conj() @ u_mat @ c_ip.T
