import numpy as np
from scipy import integrate
import plasmapy.formulary as pp_form
import astropy.units as units
from .constants import *
import warnings

#!Maybe not the best idea, but fine for now...
warnings.filterwarnings("ignore")


class DaligaultGamma:
    """
    class DaligaultGamma generates all necessary parameters and outputs for analysis
    """

    def __init__(
        self, E, Te, n, theta_cutoff=1, q_kF_list=np.linspace(1e-3, 10, 1000)
    ) -> None:
        """Function returns all parameters

        Args:
            E (array): Energy array in eV
            Te (float): electronic temperature in eV
            theta (float): degeneracy parameter
            n (float): carrier density of material (1/m^3)
            theta_cutoff (float): Te (eV) cutoff in which we go from quantum to classical chemical potential
            q_kF_list (array): array of q/kF points used for v(q) and G(q/kF), defaults to np.linspace(1e-3, 10, 1000) 1/m
        """
        self.E = E
        self.Te = Te
        self.theta = pp_form.quantum_theta(Te * units.eV, n).value
        self.n = n
        self.theta_cutoff = theta_cutoff
        self.q_kF_list = q_kF_list

    #!---------------------------------------------------------------
    # * Start of density- and theta-dependent parameter section

    def get_kF(n):
        return (3 * np.pi**2 * n) ** (1 / 3)

    def get_T_from_theta(n, theta):
        Ef = pp_form.quantum.Ef_(n * units.m**-3).value
        return theta * Ef / constants.kB

    def get_mu_quantum(n, theta):
        T = get_T_from_theta(n, theta)
        Ef = pp_form.quantum.Ef_(n * units.m**-3).value
        return Ef * (1 - 1 / 3 * ((np.pi * constants.kB * T) / (2 * Ef)) ** 2)

    # Eq. 5.26 on pp. 40 in Fetter and Walecka
    def get_mu_classical(n, theta):
        T = get_T_from_theta(n, theta)
        mu = (
            constants.kB
            * T
            * np.log(n / 2 * pp_form.quantum.lambdaDB_th_(T * units.K).value ** 3)
        )
        return mu

    def get_mu(n, theta, theta_cutoff=1):
        if theta < theta_cutoff:
            mu = get_mu_quantum(n, theta)
        elif theta >= theta_cutoff:
            mu = get_mu_classical(n, theta)
        return mu

    def get_omega_p(n):
        return pp_form.plasma_frequency(n * units.m**-3, particle="e-").value

    # * End of density-dependent parameter section
    #!---------------------------------------------------------------

    #!---------------------------------------------------------------
    # * Start of Eq. (4) section

    def get_u_integral(u, k_kF, q_kF, theta, sq_mu):
        integral_first_term = (1 / (1 - np.exp(((2 * q_kF) / theta) * u))) * (
            1
            / (
                1
                + np.exp((-(2 * q_kF) / theta) * u)
                * np.exp((1 / theta) * (sq_mu - k_kF**2))
            )
        )
        integral_second_term = np.log(
            (1 + np.exp((1 / theta) * (sq_mu - (u + q_kF / 2) ** 2)))
            / (1 + np.exp((1 / theta) * (sq_mu - (-u + q_kF / 2) ** 2)))
        )

        return integral_first_term * integral_second_term

    def get_G(k_kF, q_kF, theta, sq_mu, F_12, F_n12):

        integral, error = integrate.quad(
            get_u_integral,
            q_kF / 2 - k_kF,
            q_kF / 2 + k_kF,
            args=(k_kF, q_kF, theta, sq_mu),
        )

        G = (
            (9 * np.pi / (16 * np.sqrt(2)))
            * (theta ** (5 / 2))
            * (F_12 / F_n12) ** (3 / 2)
            * (q_kF / k_kF)
            * integral
        )

        return G, error

    # * End of Eq. (4) section
    #!---------------------------------------------------------------

    #!---------------------------------------------------------------
    # * Start of v(q) section

    def get_k_sc(n, theta, F_12, F_n12):
        T = get_T_from_theta(n, theta)
        k_DH = (
            1 / pp_form.lengths.Debye_length(T * units.K, n * units.m**-3).value
        )  # 1/m
        k_sc = k_DH * np.sqrt(F_n12 / F_12)  # 1/m

        return k_sc

    def get_v(q, k_sc):
        v = 1 / (q**2 + k_sc**2)
        return v

    # * End of v(q) section
    #!---------------------------------------------------------------

    #!---------------------------------------------------------------
    # * Start of Eq. (3) section

    def get_q_integrand(q, q_kF, k_kF, k_sc, theta, sq_mu, F_12, F_n12):
        v2 = np.abs(get_v(q, k_sc)) ** 2
        G = get_G(k_kF, q_kF, theta, sq_mu, F_12, F_n12)[0]

        q_integrand = v2 * G
        return q_integrand

    def get_Eq_3(q_list, kF, k_kF, k_sc, theta, sq_mu, F_12, F_n12):

        q_integrand = np.array(
            [
                get_q_integrand(q, q / kF, k_kF, k_sc, theta, sq_mu, F_12, F_n12)
                for q in q_list
            ]
        )
        q_integral = integrate.cumtrapz(q_integrand, q_list, initial=0)[-1]

        return k_sc**3 * q_integral

    def get_gamma_omega_p(q_kF_list, k_kF_list, n, theta, theta_cutoff):
        print(
            f"n = {n:.2e} 1/m^3, theta = {theta:.2f}, Te={get_T_from_theta(n, theta)/constants.eV_to_K:.2f}"
        )
        # * Assign values independent of k and q
        kF = get_kF(n)
        mu = get_mu(n, theta, theta_cutoff)
        sq_mu = mu / pp_form.quantum.Ef_(n * units.m**-3).value
        F_12 = np.real(pp_form.mathematics.Fermi_integral(sq_mu / theta, 1 / 2))
        F_n12 = np.real(pp_form.mathematics.Fermi_integral(sq_mu / theta, -1 / 2))
        k_sc = get_k_sc(n, theta, F_12, F_n12)
        q_list = q_kF_list * get_kF(n)
        k_list = k_kF_list * get_kF(n)

        gamma_omega_p = np.array(
            [
                get_Eq_3(q_list, kF, k / kF, k_sc, theta, sq_mu, F_12, F_n12)
                for k in k_list
            ]
        )

        return gamma_omega_p

    #!---------------------------------------------------------------
    # * Start of Eq. (3) section

    def get_gamma(self):
        """
        Function calculates Daligault gamma
        Returns:
            E_gamma (2D array): Energy-mu (eV) and Daligault scattering rate (1/fs)
        """

        k_kF_list = (
            np.sqrt(2 * constants.m_e * self.E * constants.eV_to_J) / constants.hbar
        ) / get_kF(self.n)

        gamma_list = (
            np.array(
                get_gamma_omega_p(
                    self.q_kF_list, k_kF_list, self.n, self.theta, self.theta_cutoff
                )
            )
            * get_omega_p(self.n)
            * 1e-15
        )
        # 1/fs
        E_mu_list = (
            self.E - get_mu(self.n, self.theta, self.theta_cutoff) / constants.eV_to_J
        )

        E_gamma = np.array([E_mu_list, gamma_list])  # eV and 1/fs
        return E_gamma
