"""Module for evaluating response properties from Lindhard function."""

from typing import Tuple
import math

import numpy as np

from tbplas.base.consts import NM, EPSILON0
from tbplas.base.lattice import gen_reciprocal_vectors, frac2cart, cart2frac
from tbplas.diag.base import DiagSolver
from tbplas.diag.io import load_dyn_pol, load_ac_cond


__all__ = ["Lindhard"]


class Lindhard(DiagSolver):
    """
    Lindhard function calculator.

    NOTES
    -----
    1. Units

    'Lindhard' class uses eV for energies, elementary charge of electron for
    charges and nm for lengths. So the unit for dynamic polarization is
    1 / (eV * nm**2) in 2d case and 1 / (eV * nm**3) in 3d case. Relative
    dielectric function is dimensionless. The unit for AC conductivity is
    e**2/hbar in 2d case and e**2/(hbar*nm) in 3d case.

    2. Coulomb potential and dielectric function

    The Coulomb potential in Lindhard class takes the following form:
        V(r) = 1 / epsilon * e**2 / r
             = 1 / (epsilon_0 * epsilon_r) * e**2 / r
    where epsilon is the ABSOLUTE dielectric electric constant of background.

    Now we evaluate the absolute dielectric constant of vacuum in the units of
    eV/nm/q. Suppose that we have two electrons separated at the distance of
    1nm, then their potential energy in SI units is:
        V = 1 / (4 * pi * epsilon_0) * q**2 / 1e-9 [Joule]
          = 1 / (4 * pi * epsilon_0) * q**2 / 1e-9 * 1 / q [eV]
          = q * 1e9 / (4 * pi * epsilon_0) [eV]
    The counterpart in eV/nm/q units is:
        V = 1 / epsilon_0_arb [eV]
    So we have:
        1 / epsilon_0_arb = q * 1e9 / (4 * pi * epsilon_0)
    where the right part are the values in SI. Finally, we have:
        epsilon_0_arb = 0.6944615417149689
    in eV/nm/q units. That's EPSILON0 in constants module.

    3. Fourier transform of Coulomb potential

    The 3D Fourier transform of Coulomb potential V(r) as defined in section 2
    takes the following form:
        V(q) = 4 * pi * e**2 / (epsilon_0 * epsilon_r * q**2)
    For the derivation, see pp. 19 of Many-Particle Physics by Gerald D. Mahan.

    The 2D Fourier transform takes the form:
        V(q) = 2 * pi * e**2 / (epsilon_0 * epsilon_r * q)
    See the following url for the derivation:
    https://math.stackexchange.com/questions/3627267/fourier-transform-of-the-2d-coulomb-potential

    4. System dimension

    'Lindhard' class deals with system dimension in two approaches. The first
    approach is to treat all systems as 3-dimensional. Supercell technique is
    required in this approach, with vacuum layers added on non-periodic
    directions. Also, the component(s) of kmesh_size should be set to 1
    accordingly.

    The second approach utilizes dimension-specific formula whenever possible.
    For now, only 2-dimensional case has been implemented. This approach
    requires that the system should be periodic in xOy plane, i.e. the
    non-periodic direction should be along 'c' axis. Otherwise, the results
    will be wrong.

    The first approach suffers from the problem that dynamic polarization and
    AC conductivity scale inversely proportional to the product of supercell
    lengths, i.e. |c| in 2d case and |a|*|b| in 1d case. This is due to the
    elementary volume in reciprocal space (dnk) in the summation in Lindhard
    function. On the contrary, the second approach has no such issue. If the
    supercell lengths of non-periodic directions are set to 1 nm, then the first
    approach yields the same results as the second approach.

    For the dielectric function, the situation is more complicated. From the
    equation
        epsilon(q) = 1 - V(q) * dyn_pol(q)
    we can see it is also affected by the Coulomb potential V(q). Since
        V(q) = 4 * pi * e**2 / (epsilon_0 * epsilon_r * q**2) (3d)
    and
        V(q) = 2 * pi * e**2 / (epsilon_0 * epsilon_r * q) (2d)
    the influence of system dimension is q-dependent. Setting supercell lengths
    to 1 nm will NOT produce the same result as the second approach.

    5. Dielectric function for q = 0

    The dielectric function for q = 0 cannot be evaluated from dynamic
    polarizability as v(q) diverges. Instead, we calculate it from ac
    conductivity. We have two formulae. The first one is eqn. (63)
        epsilon = 1 + 1j * 4 * pi / omega * sigma
    in https://www.openmx-square.org/tech_notes/Dielectric_Function_YTL.pdf,
    while the second one is eqn. (1.10)
        epsilon = 1 + 1j / (epsilon_0 * omega) * sigma
    in Plasmonics: Fundamentals and Applications by Stefan A. Maier. The former
    is in Gaussian units, while the latter is in SI units. Unfortunately, we
    cannot apply them directly.

    As we take the Coulomb potential to be
            V(r) = 1 / epsilon * e**2 / r
                 = 1 / (epsilon_0 * epsilon_r) * e**2 / r
    there should be a 4pi factor as in Gaussian units. However, epsilon_0 is not
    1 as we are not using cm/g/s as the fundamental units. The actual formula is
        epsilon = 1 + 1j * 4 * pi / (epsilon_0 * omega) * sigma
    which can be regarded as a mixture of the formulae in Gaussian and SI units.

    The unit of epsilon follows:
        [epsilon] = [sigma] / [epsilon_0 * omega]
                  = e**2/(hbar*nm) / [e**2/(eV*nm) * eV/hbar]
                  = e**2/(hbar*nm) / [e**2/(hbar*nm)]
                  = 1
    However, this equation holds for 3d systems only. For 2d systems there will
    be a remaining unit of nm.
    """
    def _check_k_grid_dim(self) -> None:
        """Check k_grid_size and dimension."""
        k_grid_size = self.config.k_grid_size
        dimension = self.config.dimension
        assert len(k_grid_size) == 3, "Length of k_grid_size should be 3"
        for i in k_grid_size:
            assert i >= 1, f"Dimension of k_grid_size should >= 1"
        assert dimension in (2, 3), f"Unsupported dimension: {dimension}"
        if dimension == 2:
            assert k_grid_size[2] == 1, ("2d specific algorithms require"
                                         " kmesh_size[2] == 1")

    def k_frac2cart(self,
                    frac_coord: np.ndarray,
                    unit: float = NM) -> np.ndarray:
        """
        Convert FRACTIONAL coordinates of k-points to CARTESIAN coordinates
        in 1/unit.
        :param frac_coord: (num_kpt, 3) float64 array
            FRACTIONAL coordinates of k-points
        :param unit: scaling factor from unit to NANOMETER,
            e.g. unit=0.1 for ANGSTROM
        :return: (num_kpt, 3) float64 array
            CARTESIAN coordinates of k-points in 1/unit
        """
        if not isinstance(frac_coord, np.ndarray):
            frac_coord = np.array(frac_coord, dtype=np.float64)
        recip_lat_vec = gen_reciprocal_vectors(self._model.lat_vec)
        cart_coord_nm = frac2cart(recip_lat_vec, frac_coord)
        cart_coord = cart_coord_nm * unit
        return cart_coord

    def k_cart2frac(self,
                    cart_coord: np.ndarray,
                    unit: float = NM) -> np.ndarray:
        """
        Convert CARTESIAN coordinates of k-points in 1/unit to FRACTIONAL
        coordinates.
        :param cart_coord: (num_kpt, 3) float64 array
            CARTESIAN coordinates in 1/unit
        :param unit: scaling factor from unit to NANOMETER,
            e.g. unit=0.1 for ANGSTROM
        :return: (num_kpt, 3) float64 array
            FRACTIONAL coordinates of k-points
        """
        if not isinstance(cart_coord, np.ndarray):
            cart_coord = np.array(cart_coord, dtype=np.float64)
        recip_lat_vec = gen_reciprocal_vectors(self._model.lat_vec)
        cart_coord_nm = cart_coord / unit
        frac_coord = cart2frac(recip_lat_vec, cart_coord_nm)
        return frac_coord

    def calc_dyn_pol(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate dynamic polarizability for arbitrary q-points.
        :return: (omegas, dyn_pol)
            omegas: (num_omega,) float64 array
            energies in eV
            dyn_pol: (num_qpt, num_omega) complex128 array
            dynamic polarizability in 1/(eV*nm**2) or 1/(eV*nm**3)
        """
        if __debug__:
            self._check_k_grid_dim()
        self._pre_run("lindhard")
        self._cpp_solver.calc_dyn_pol()
        return load_dyn_pol(self.config.prefix)

    def _get_vq(self, q_point: np.ndarray) -> float:
        """
        Get Coulomb interaction in momentum space for given q-point.
        :param q_point: (3, ) float64 array
            CARTESIAN coordinate of q-point in 1/NM
        :return: Coulomb interaction in momentum space in eV
        """
        factor = 1 / (EPSILON0 * self.config.back_epsilon)
        q_norm = np.linalg.norm(q_point)
        if self.config.dimension == 2:
            vq = factor * 2 * math.pi / q_norm
        else:
            vq = factor * 4 * math.pi / q_norm**2
        return vq

    def calc_epsilon(self, dyn_pol: np.ndarray) -> np.ndarray:
        """
        Calculate dielectric function from dynamic polarization.

        :param dyn_pol: (num_qpt, num_omega) complex128 array
            dynamic polarization from calc_dyn_pol
        :return: (num_qpt, num_omega) complex128 array
            relative dielectric function
        """
        if __debug__:
            self._check_k_grid_dim()
        num_qpt = dyn_pol.shape[0]
        num_omega = dyn_pol.shape[1]
        assert num_qpt == self.config.q_points.shape[0]
        epsilon = np.zeros((num_qpt, num_omega), dtype=np.complex128)
        for i_q, q_point in enumerate(self.config.q_points):
            vq = self._get_vq(q_point)
            epsilon[i_q] = 1.0 - vq * dyn_pol[i_q]
        return epsilon

    def calc_ac_cond(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate AC conductivity using Kubo-Greenwood formula.

        Reference: section 12.2 of Wannier90 user guide.
        NOTE: there is not such g_s factor in the reference.

        :return: (omegas, ac_cond)
            omegas: (num_omega,) float64 array
            energies in eV
            ac_cond: (num_omega,) complex128 array
            AC conductivity in e**2/(hbar*nm) in 3d case and e**2/hbar
            in 2d case
        """
        if __debug__:
            self._check_k_grid_dim()
            assert self.config.ac_cond_component in range(1, 10)
        self._pre_run("lindhard")
        self._cpp_solver.calc_ac_cond()
        return load_ac_cond(self.config.prefix)

    def calc_epsilon_q0(self,
                        omegas: np.ndarray,
                        ac_cond: np.ndarray) -> np.ndarray:
        """
        Calculate dielectric function from AC conductivity for q=0.
        :param omegas: (num_omega,) float64 array
            energies in eV
        :param ac_cond: (num_omega,) complex128 array
            AC conductivity in e**2/(hbar*nm) in 3d case
        :return: (num_omega,) complex128 array
            relative dielectric function
        """
        if __debug__:
            self._check_k_grid_dim()
        assert self.config.dimension == 3, (f"Unsupported dimension:"
                                            f" {self.config.dimension}")
        prefactor = 4 * math.pi / (self.config.back_epsilon * EPSILON0)
        return 1 + 1j * prefactor * ac_cond / omegas
