"""Functions and classes for analyzing correlation functions."""

from math import cos, sin, exp, pi
from typing import Callable, Tuple

import numpy as np
from scipy.signal import hilbert
try:
    from scipy.integrate import trapz
except ImportError:
    from scipy.integrate import trapezoid as trapz

from tbplas.base.consts import EPSILON0, KB_EVK
from tbplas.base.parallel import MPIEnv
from tbplas.tbpm.config import JobInfo
from tbplas.tbpm.solver import get_qe_indices


__all__ = ["window_hanning", "window_exp", "window_exp_ten", "Analyzer"]


def window_hanning(i: int, tnr: int) -> float:
    """
    Hanning window function.
    :param i: summation index
    :param tnr: total length of summation
    :return: Hanning window value
    """
    return 0.5 * (1 + cos(pi * i / tnr))


def window_exp(i: int, tnr: int) -> float:
    """
    Exponential window function.
    :param i: summation index
    :param tnr: total length of summation
    :return: exponential window value
    """
    return exp(-2. * (i / tnr)**2)


def window_exp_ten(i: int, tnr: int) -> float:
    """
    Exponential window function with base 10.
    :param i: summation index
    :param tnr: total length of summation
    :return: exponential window value
    """
    power = -2 * (1. * i / tnr)**2
    return 10.**power


class Analyzer:
    """
    Class for analyzing correlation functions.

    Attributes
    ----------
    _info: 'JobInfo' instance
        parameters controlling TBPM calculation
    _mpi_env: 'MPIEnv' instance
        mpi parallel environment
    """
    def __init__(self, info_name: str, echo_details: bool = False) -> None:
        """
        :param info_name: file name of *_info.dat
        :param echo_details: whether to output parallelization details
        """
        self._info = JobInfo(info_name)
        self._mpi_env = MPIEnv(echo_details=echo_details)

    def calc_dos(self,
                 corr_dos: np.ndarray,
                 window: Callable = window_hanning) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate DOS from correlation function.

        Reference: eqn. 16-17 of feature article.

        The unit of dos follows:
            [dos] = [C_DOS] [dt] = hbar / eV
        So possibly the formula misses a hbar on the denominator.
        Anyway, the DOS is correct since it is explicitly normalized to 1.

        :param corr_dos: (num_steps+1, ) complex128 array
            dimensionless DOS correlation function
        :param window: window function for integral
        :return: (energies, dos)
            energies: (2*num_steps,) float64 array
            energies in eV
            dos: (2*num_steps,) float64 array
            DOS in 1/eV
        """
        self._info.load()
        assert self._info.num_spin in (1, 2)

        # Get parameters
        tnr = self._info.num_time_steps
        en_range = self._info.energy_range
        en_step = 0.5 * en_range / tnr

        # Allocate working arrays
        energies = np.array([0.5 * i * en_range / tnr - en_range / 2.
                             for i in range(tnr * 2)], dtype=float)
        dos = np.zeros(tnr * 2, dtype=float)

        if self._mpi_env.is_master:
            # Get negative time correlation
            corr_neg_time = np.empty(tnr * 2, dtype=complex)
            corr_neg_time[tnr - 1] = corr_dos.item(0)
            corr_neg_time[2 * tnr - 1] = window(tnr - 1, tnr) * corr_dos.item(tnr)
            for i in range(tnr - 1):
                corr_neg_time[tnr + i] = window(i, tnr) * corr_dos.item(i + 1)
                corr_neg_time[tnr-i-2] = window(i, tnr) * corr_dos.item(i + 1).conjugate()

            # Fourier transform
            corr_fft = np.fft.ifft(corr_neg_time)
            for i in range(tnr):
                dos[i + tnr] = abs(corr_fft.item(i))
            for i in range(tnr, 2 * tnr):
                dos[i - tnr] = abs(corr_fft.item(i))

            # Normalise and correct for spin
            dos = dos / (np.sum(dos) * en_step)
            dos *= self._info.num_spin
        else:
            pass
        self._mpi_env.bcast(dos)
        return energies, dos

    def calc_ac_cond(self,
                     corr_ac: np.ndarray,
                     window: Callable = window_exp) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate optical (AC) conductivity from correlation function.

        Reference: eqn. 300-301 of graphene note.

        The unit of AC conductivity in 2d case follows:
        [sigma] = [1/(hbar * omega * A)] * [j^2] * [dt]
                = 1/(eV*nm^2) * e^2/hbar^2 * (eV)^2 * nm^2 * hbar/eV
                = e^2/hbar
        which is consistent with the results from Lindhard function.

        The reason for num_orbitals in the prefactor is that every electron
        contribute freely to the conductivity, and we have to take the number
        of electrons into consideration. See eqn. 222-223 of the note for more
        details.

        :param corr_ac: (4, num_steps) complex128 array
            AC correlation function along xx, yx, xy, yy directions
            Unit is e^2/hbar^2 * (eV)^2 * nm^2.
        :param window: window function for integral
        :return: (omegas, ac_cond)
            omegas: (num_steps,) float64 array
            frequencies in eV
            ac_cond: (4, num_steps) complex128 array
            ac conductivity values corresponding to omegas for 4 directions
            (xx, yx, xy, yy, respectively)
            The unit is e^2/(hbar*nm) in 3d case and e^2/hbar in 2d case.
        """
        self._info.load()
        assert self._info.num_spin in (1, 2)
        assert self._info.dimension in (2, 3)

        # Get parameters
        tnr = self._info.num_time_steps
        en_range = self._info.energy_range
        t_step = pi / en_range
        beta = 1.0 / (KB_EVK * self._info.temperature)
        if self._info.dimension == 2:
            ac_prefactor = self._info.num_orbitals \
                           / self._info.lattice_area
        else:
            ac_prefactor = self._info.num_orbitals \
                           / self._info.lattice_volume

        # Allocate working arrays
        omegas = np.array([i * en_range / tnr for i in range(tnr)],
                          dtype=float)
        ac_cond = np.zeros((4, tnr), dtype=complex)

        if self._mpi_env.is_master:
            # Get real part of AC conductivity
            ac_real = np.zeros((4, tnr), dtype=float)
            for j in range(4):
                for i in range(tnr):
                    omega = omegas.item(i)
                    acv = 0.
                    for k in range(tnr):
                        acv += 2. * window(k + 1, tnr) \
                            * sin(omega * k * t_step) * corr_ac.item(j, k).imag
                    if omega == 0.:
                        acv = 0.
                    else:
                        acv = ac_prefactor * t_step * acv \
                            * (exp(-beta * omega) - 1) / omega
                    ac_real[j, i] = acv

            # Get imaginary part of AC conductivity via Kramers-Kronig relations
            # (Hilbert transformation).
            ac_imag = np.zeros((4, tnr), dtype=float)
            for j in range(4):
                sigma = np.zeros(2 * tnr, dtype=float)
                for i in range(tnr):
                    sigma[tnr + i] = ac_real.item(j, i)
                    sigma[tnr - i] = ac_real.item(j, i)
                ac_imag[j, :] = np.imag(hilbert(sigma))[tnr:2 * tnr]
            ac_cond = ac_real + 1j * ac_imag

            # Correct for spin
            ac_cond *= self._info.num_spin
        else:
            pass
        self._mpi_env.bcast(ac_cond)
        return omegas, ac_cond

    def calc_dyn_pol(self,
                     corr_dyn_pol: np.ndarray,
                     window: Callable = window_exp_ten) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        Calculate dynamical polarization from correlation function.

        Reference:
        https://journals.aps.org/prb/abstract/10.1103/PhysRevB.84.035439

        The unit of dp in 2d case follows:
        [dp] = [1/A] * [C_DP] * [dt]
             = 1/nm^2 * 1 * hbar/eV
             = hbar/(eV*nm^2)
        which is inconsistent with the output of Lindhard! So, possibly the
        formula misses a hbar on the denominator.

        :param corr_dyn_pol: (num_qpt, num_steps) float64 array
            dimensionless dynamical polarization correlation function.
        :param window: window function for integral
        :return: (q_points, omegas, dyn_pol)
            q_points: (num_qpt, 3) float64 array
            Cartesian coordinates of q-points in 1/nm
            omegas: (num_steps,) float64 array
            frequencies in eV
            dyn_pol: (num_qpt, num_steps) complex128 array
            dynamical polarization values corresponding to q-points and omegas
            The unit is 1/(eV*nm^2) or 1/(eV*nm^3) depending on the dimension.
        """
        self._info.load()
        assert self._info.num_spin in (1, 2)
        assert self._info.dimension in (2, 3)

        # Get parameters
        tnr = self._info.num_time_steps
        en_range = self._info.energy_range
        t_step = pi / en_range
        q_points = np.array(self._info.dyn_pol_q_points, dtype=float)
        if self._info.dimension == 2:
            dyn_pol_prefactor = 2. * self._info.num_orbitals \
                                / self._info.lattice_area
        else:
            dyn_pol_prefactor = 2. * self._info.num_orbitals \
                                / self._info.lattice_volume

        # Allocate working arrays
        n_q_points = len(q_points)
        n_omegas = tnr
        omegas = np.array([i * en_range / tnr for i in range(tnr)],
                          dtype=float)
        dyn_pol = np.zeros((n_q_points, n_omegas), dtype=complex)

        if self._mpi_env.is_master:
            for i_q in range(n_q_points):
                for i in range(n_omegas):
                    omega = omegas.item(i)
                    dpv = 0.0j
                    for k in range(tnr):
                        phi = k * t_step * omega
                        dpv += window(k + 1, tnr) * corr_dyn_pol.item(i_q, k) \
                            * (cos(phi) + 1j * sin(phi))
                    dyn_pol[i_q, i] = dyn_pol_prefactor * t_step * dpv

            # correct for spin
            dyn_pol *= self._info.num_spin
        else:
            pass
        self._mpi_env.bcast(dyn_pol)
        return q_points, omegas, dyn_pol

    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
        """
        assert self._info.dimension in (2, 3)
        factor = self._info.dyn_pol_coulomb_factor
        back_epsilon = self._info.dyn_pol_back_epsilon
        factor /= (EPSILON0 * back_epsilon)
        q_norm = np.linalg.norm(q_point)
        if q_norm == 0.0:
            vq = 0.0
        else:
            if self._info.dimension == 2:
                vq = factor * 2 * pi / q_norm
            else:
                vq = factor * 4 * pi / q_norm**2
        return vq

    def calc_epsilon(self, dyn_pol: np.ndarray) -> np.ndarray:
        """
        Calculate dielectric function from dynamical polarization.
        :param dyn_pol: (num_qpt, num_steps) complex128 array
            dynamical polarization
        :return: (num_qpt, num_steps) complex128 array
            dielectric function
        """
        # Get parameters
        self._info.load()
        tnr = self._info.num_time_steps
        q_points = self._info.dyn_pol_q_points
        n_q_points = len(q_points)
        n_omegas = tnr
        epsilon = np.zeros((n_q_points, n_omegas), dtype=complex)

        if self._mpi_env.is_master:
            for i, q_point in enumerate(q_points):
                vq = self._get_vq(np.array(q_point))
                epsilon[i] = 1 - vq * dyn_pol[i]
        else:
            pass
        self._mpi_env.bcast(epsilon)
        return epsilon

    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_steps,) float64 array
            energies in eV
        :param ac_cond: (4, num_steps) complex128 array
            AC conductivity in e**2/(hbar*nm) in 3d case along xx, yx, xy, yy
            directions
        :return: (4, num_steps) complex128 array
            relative dielectric function
        """
        self._info.load()
        assert self._info.dimension == 3
        back_epsilon = self._info.dyn_pol_back_epsilon
        prefactor = 4 * pi / (back_epsilon * EPSILON0)
        return 1 + 1j * prefactor * ac_cond / omegas

    def calc_dc_cond(self,
                     corr_dos: np.ndarray,
                     corr_dc: np.ndarray,
                     window_dos: Callable = window_hanning,
                     window_dc: Callable = window_exp) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate electronic (DC) conductivity at zero temperature from its
        correlation function.

        Reference: eqn. 381 of graphene note.

        The unit of DC conductivity in 2d case follows:
        [sigma] = [D/A] * [j^2] * [dt]
                = 1/(eV*nm^2) * e^2/hbar^2 * (eV)^2 * nm^2 * hbar/eV
                = e^2/hbar
        which is consistent with the result of Lindhard function.

        NOTE: the xy and yx components of conductivity are not accurate. So
        they will not be evaluated.

        NOTE: Here we need to call analyze_corr_dos to obtain DOS, which is
        intended to analyze the result of calc_corr_dos by design. As in
        the c++ extension, the results of calc_corr_dos and calc_corr_ldos
        have the length of num_steps+1, while that of calc_corr_dc has
        length of num_steps. This is due to incomplete update of the
        source code. tbpm_dos and tbpm_ldos have been update, while other
        subroutines are not. So here we need to insert 1.0 to the head of
        corr_dos by calc_corr_dc before calling analyze_corr_dos.

        :param corr_dos: corr_dos: (num_steps,) complex128 array
            dimensionless DOS correlation function
        :param corr_dc: (2, num_eng, num_steps_dc) complex128 array
            DC conductivity correlation function in e^2/hbar^2 * (eV)^2 * nm^2
        :param window_dos: window function for DOS integral
        :param window_dc: window function for DC integral
        :return: (energies, dc)
            energies: (num_eng,) float64 array
            energies in eV
            dc: (2, num_eng) float64 array
            dc conductivity values for xx and yy directions in the same unit
            as ac conductivity
        """
        self._info.load()
        assert self._info.num_spin in (1, 2)
        assert self._info.dimension in (2, 3)

        # Get DOS
        corr_dos = np.insert(corr_dos, 0, 1.0)
        energies_dos, dos = self.calc_dos(corr_dos, window_dos)
        energies_dos = np.array(energies_dos)
        dos = np.array(dos)

        # Get parameters
        tnr_dc = self._info.dc_cond_num_time_steps
        en_range = self._info.energy_range
        t_step = 2 * pi / en_range
        qe_indices = get_qe_indices(self._info, energies_dos)

        n_energies = len(qe_indices)
        energies = energies_dos[qe_indices]
        if self._info.dimension == 2:
            dc_prefactor = self._info.num_orbitals / \
                           self._info.lattice_area
        else:
            dc_prefactor = self._info.num_orbitals / \
                           self._info.lattice_volume
        dc = np.zeros((2, n_energies))

        if self._mpi_env.is_master:
            for i in range(2):
                for j in range(n_energies):
                    en = energies.item(j)
                    dos_val = dos.item(qe_indices.item(j))
                    dc_val = 0.
                    for k in range(tnr_dc):
                        w = window_dc(k + 1, tnr_dc)
                        phi = k * t_step * en
                        c_exp = cos(phi) - 1j * sin(phi)
                        add_dcv = w * (c_exp * corr_dc.item(i, j, k)).real
                        dc_val += add_dcv
                    dc[i, j] = dc_prefactor * t_step * dos_val * dc_val
            # correct for spin
            dc *= self._info.num_spin
        else:
            pass
        self._mpi_env.bcast(dc)
        return energies, dc

    def calc_diff_coeff(self,
                        corr_dc: np.ndarray,
                        window_dc: Callable = window_exp) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate diffusion coefficient form DC correlation function.

        Reference: eqn. 43-44 of feature article.

        The unit of diff_coeff follows:
        [diff_coeff] = [1/e^2] * [j^2] * [dt]
                     = 1/e^2 * e^2/hbar^2 * (eV)^2 * nm^2 * hbar/eV
                     = eV*nm^2/hbar
        which does not depend on system dimension.

        TODO: consider spin degeneracy?

        :param corr_dc: (2, num_eng, num_steps_dc) complex128 array
            DC conductivity correlation function in e^2/hbar^2 * (eV)^2 * nm^2
        :param window_dc: window function for DC integral
        :return: (time, diff_coeff)
            time: (num_steps,) float64 array
            time for diffusion coefficient in hbar/eV
            diff_coeff: (2, num_eng, num_steps) complex128 array
            diffusion coefficient in nm^2/(hbar/eV)
        """
        self._info.load()
        assert self._info.num_spin in (1, 2)
        assert self._info.dimension in (2, 3)

        # Get parameters
        tnr = self._info.num_time_steps
        en_range = self._info.energy_range
        t_step = 2 * pi / en_range
        energies = np.array([0.5 * i * en_range / tnr - en_range / 2.
                             for i in range(tnr * 2)], dtype=float)
        qe_indices = get_qe_indices(self._info, energies)
        n_energies = len(qe_indices)
        energies = energies[qe_indices]
        time = np.linspace(0, tnr - 1, tnr) * t_step
        diff_coeff = np.zeros((2, n_energies, tnr))

        if self._mpi_env.is_master:
            for i in range(2):
                for j in range(n_energies):
                    en = energies.item(j)
                    temp = np.zeros(tnr)
                    for k in range(tnr):
                        w = window_dc(k + 1, tnr)
                        phi = k * t_step * en
                        c_exp = cos(phi) - 1j * sin(phi)
                        temp[k] = w * (c_exp * corr_dc.item(i, j, k)).real
                    for k2 in range(tnr):
                        diff_coeff[i, j, k2] = trapz(temp[:k2 + 1], time[:k2 + 1])
        else:
            pass
        self._mpi_env.bcast(diff_coeff)
        return time, diff_coeff

    def calc_dc_cond_t(self,
                       corr_dos: np.ndarray,
                       corr_dc: np.ndarray,
                       window_dos: Callable = window_hanning,
                       window_dc: Callable = window_exp) -> Tuple[np.ndarray, np.ndarray]:
        """
        Similar as calc_dc_cond, but keeps the integral of correlation function
        up to each time step.
        :param corr_dos: (num_steps,) complex128 array
            dimensionless DOS correlation function
        :param corr_dc: (2, num_eng, num_steps_dc) complex128 array
            DC conductivity correlation function in e^2/hbar^2 * (eV)^2 * nm^2
        :param window_dos: window function for DOS integral
        :param window_dc: window function for DC integral
        :return: (energies, dc)
            energies: (num_eng,) float64 array
            energies in eV
            dc: (2, num_eng, num_steps_dc) float64 array
            dc conductivity values integrated up to each time step for xx and yy
            directions in the same unit as ac conductivity
        """
        self._info.load()
        assert self._info.num_spin in (1, 2)
        assert self._info.dimension in (2, 3)

        # Get DOS
        corr_dos = np.insert(corr_dos, 0, 1.0)
        energies_dos, dos = self.calc_dos(corr_dos, window_dos)
        energies_dos = np.array(energies_dos)
        dos = np.array(dos)

        # Get parameters
        tnr_dc = self._info.dc_cond_num_time_steps
        en_range = self._info.energy_range
        t_step = 2 * pi / en_range
        qe_indices = get_qe_indices(self._info, energies_dos)

        n_energies = len(qe_indices)
        energies = energies_dos[qe_indices]
        if self._info.dimension == 2:
            dc_prefactor = self._info.num_orbitals / \
                           self._info.lattice_area
        else:
            dc_prefactor = self._info.num_orbitals / \
                           self._info.lattice_volume
        dc = np.zeros((2, n_energies, tnr_dc))

        if self._mpi_env.is_master:
            for i in range(2):
                for j in range(n_energies):
                    en = energies.item(j)
                    dos_val = dos.item(qe_indices.item(j))
                    dc_val = 0.
                    for k in range(tnr_dc):
                        w = window_dc(k + 1, tnr_dc)
                        phi = k * t_step * en
                        c_exp = cos(phi) - 1j * sin(phi)
                        add_dcv = w * (c_exp * corr_dc.item(i, j, k)).real
                        dc_val += add_dcv
                        dc[i, j, k] = dc_prefactor * t_step * dos_val * dc_val
            # correct for spin
            dc *= self._info.num_spin
        else:
            pass
        self._mpi_env.bcast(dc)
        return energies, dc

    def calc_bands(self, corr_bands: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate band structure from correlation function.
        :param corr_bands: (num_kpt, num_steps+1) complex128 array
            correlation functions for each k-point
        :return: (energy, bands)
            energy: (num_eng,) float64 array
            energies in eV
            bands: (num_kpt, num_eng) complex128 array
            weight function (DOS) for each k-point
        """
        energy_all, dos_all = [], []
        for corr_dos in corr_bands:
            eng, dos = self.calc_dos(corr_dos)
            energy_all.append(eng)
            dos_all.append(dos)
        dos_all = np.array(dos_all)
        return energy_all[0], dos_all

    @property
    def is_master(self) -> bool:
        """Wrapper for determine the master process for compatibility."""
        return self._mpi_env.is_master
