import numpy as np
import astropy.units as u
from astropy.constants import e, c, m_e, sigma_T
from astropy.coordinates import Distance
import matplotlib.pyplot as plt
from .. import spectra
from ..utils.conversion import mec2, B_to_cgs
from ..utils.plot import load_mpl_rc

__all__ = ["Blob"]


def init_spectrum_norm_dict(norm, spectrum_dict, norm_type="integral", V_b=None):
    """initialize a spectrum from a normalisation and a spectrum dictionary of 
    the following type
    
    .. code-block:: python

        spectrum_dict = {
            "type": "PowerLaw", 
            "parameters": {
                "p": 2.8, 
                "gamma_min": 1e2, 
                "gamma_max": 1e7
            }
        }

    Parameters
    ----------
    norm : :class:`~astropy.units.Quantity`
        normalisation of the spectrum, might be in cm-3, erg cm-3 or erg,
        if erg are specified then V_b has to be provided
    spectrum_dict : dictionary
        dictionary with the spectrum type and parameters
    norm_type : ["integral", "differential", "gamma=1"]
        normalisation type
    V_b : :class:`~astropy.units.Quantity`
        volume of the emission region, to be provided only if the normalisation
        is in erg
    """
    model = getattr(spectra, spectrum_dict["type"])

    if norm.unit in (u.Unit("erg"), u.Unit("erg cm-3")) and norm_type != "integral":
        raise NameError(
            "Normalisations different than 'integral' available only for 'spectrum_norm' in cm-3"
        )

    # check the units of the normalisation
    # cm-3 is the only one allowing more than one normalisation type
    if norm.unit == u.Unit("cm-3"):
        if norm_type == "differential":
            final_model = model(norm, **spectrum_dict["parameters"])
        elif norm_type == "gamma=1":
            final_model = model.from_norm_at_gamma_1(
                norm, **spectrum_dict["parameters"]
            )
        elif norm_type == "integral":
            final_model = model.from_normalised_density(
                norm, **spectrum_dict["parameters"]
            )

    elif norm.unit == u.Unit("erg cm-3"):
        final_model = model.from_normalised_energy_density(
            norm, **spectrum_dict["parameters"]
        )

    elif norm.unit == u.Unit("erg"):
        if V_b is None:
            raise ValueError(
                "if a normalisation in erg is provided, the volume V_b of the emission region msut be specified"
            )
        else:
            final_model = model.from_total_energy(
                norm, V_b, **spectrum_dict["parameters"]
            )

    return final_model


class Blob:
    r"""Simple spherical emission region.

    **Note:** all these quantities are defined in the comoving frame so they are actually
    primed quantities, when referring the notation in [DermerMenon2009]_.

    Parameters
    ----------
    R_b : :class:`~astropy.units.Quantity`
        radius of the blob
    z : float
        redshift of the source
    delta_D : float
        Doppler factor of the relativistic outflow
    Gamma : float
        Lorentz factor of the relativistic outflow
    B : :class:`~astropy.units.Quantity`
        magnetic field in the blob (Gauss)
    xi : float
        acceleration coefficient :math:`\xi` for first-order Fermi acceleration
        :math:`(\mathrm{d}E/\mathrm{d}t \propto v \approx c)`
        used to compute limits on the maximum Lorentz factor via
        :math:`(\mathrm{d}E/\mathrm{d}t)_{\mathrm{acc}} = \xi c E / R_L`

    spectrum_norm : :class:`~astropy.units.Quantity`
        normalisation of the electron spectra, by default can be, following 
        the notation in [DermerMenon2009]_:

            - :math:`n_{e,\,tot}`: total electrons density, in :math:`\mathrm{cm}^{-3}`
            - :math:`u_e` : total electrons energy density, in :math:`\mathrm{erg}\,\mathrm{cm}^{-3}`
            - :math:`W_e` : total energy in electrons, in :math:`\mathrm{erg}`
        
        see `spectrum_norm_type` for more details on the normalisation

    spectrum_dict : dictionary
        dictionary containing type and spectral shape information, e.g.:

        .. code-block:: python

            spectrum_dict = {
                "type": "PowerLaw", 
                "parameters": {
                    "p": 2.8, 
                    "gamma_min": 1e2, 
                    "gamma_max": 1e7
                }
            }

    spectrum_norm_type : ["integral", "differential", "gamma=1"]
        only with a normalisation in "cm-3" one can select among three types: 

        * ``"integral"``: (default) the spectrum is set such that :math:`n_{e,\,tot}` equals the value provided by ``spectrum_norm``;  
        
        * ``"differential"``: the spectrum is set such that :math:`k_e` equals the value provided by ``spectrum_norm``;    
        
        * ``"gamma=1"``: the spectrum is set such that :math:`n_e(\gamma=1)` equals the value provided by ``spectrum_norm``.
        
    gamma_size : int
        size of the array of electrons Lorentz factors
    """

    def __init__(
        self,
        R_b=1e16 * u.cm,
        z=0.069,
        delta_D=10,
        Gamma=10,
        B=1 * u.G,
        spectrum_norm=1e48 * u.Unit("erg"),
        spectrum_dict={
            "type": "PowerLaw",
            "parameters": {"p": 2.8, "gamma_min": 1e2, "gamma_max": 1e7},
        },
        spectrum_norm_type="integral",
        xi=1.0,
        gamma_size=200,
    ):
        self.R_b = R_b.to("cm")
        self.z = z
        self.d_L = Distance(z=self.z).cgs
        self.V_b = 4 / 3 * np.pi * np.power(self.R_b, 3)
        self.delta_D = delta_D
        self.Gamma = Gamma
        self.Beta = np.sqrt(1 - 1 / np.power(self.Gamma, 2))
        # viewing angle
        self.mu_s = (1 - 1 / (self.Gamma * self.delta_D)) / self.Beta
        self.theta_s = (np.arccos(self.mu_s) * u.rad).to("deg")
        self.B = B
        # B decomposed in Gaussian-cgs units
        self.B_cgs = B_to_cgs(B)
        self.spectrum_norm = spectrum_norm
        self.spectrum_norm_type = spectrum_norm_type
        self.spectrum_dict = spectrum_dict
        self.xi = xi
        # default grid of Lorentz factors for integration in the external frame
        self.gamma_to_integrate = np.logspace(1, 9, gamma_size)
        # model for the electron density
        self.set_spectrum(
            self.spectrum_norm, self.spectrum_dict, self.spectrum_norm_type, gamma_size
        )

    def set_gamma_size(self, gamma_size):
        """change size of the array of electrons Lorentz factors"""
        self.gamma_size = gamma_size
        self.gamma = np.logspace(
            np.log10(self.gamma_min), np.log10(self.gamma_max), self.gamma_size
        )
        self.gamma_to_integrate = np.logspace(1, 9, self.gamma_size)

    def set_gamma(self, gamma_min, gamma_max, gamma_size):
        """set the array of Lorentz factors to be used for integration in the 
        frame comoving with the blob"""
        self.gamma_min = gamma_min
        self.gamma_max = gamma_max
        self.gamma_size = gamma_size
        # grid of Lorentz factor for the integration in the blob comoving frame
        self.gamma = np.logspace(
            np.log10(self.gamma_min), np.log10(self.gamma_max), self.gamma_size
        )

    def set_spectrum(
        self, spectrum_norm, spectrum_dict, spectrum_norm_type, gamma_size=200
    ):
        r"""set the spectrum :math:`n_e` for the electrons accelerated in the 
        blob, reset also the array of Lorentz factor given the `gamma_min` and
        `gamma_max` in the parameters dictionary"""
        self.set_gamma(
            spectrum_dict["parameters"]["gamma_min"],
            spectrum_dict["parameters"]["gamma_max"],
            gamma_size,
        )
        self.n_e = init_spectrum_norm_dict(
            spectrum_norm, spectrum_dict, spectrum_norm_type, self.V_b
        )

    def __str__(self):
        """printable summary of the blob"""
        return (
            "* spherical emission region\n"
            + f" - R_b (radius of the blob): {self.R_b.cgs:.2e}\n"
            + f" - V_b (volume of the blob): {self.V_b.cgs:.2e}\n"
            + f" - z (source redshift): {self.z:.2f}\n"
            + f" - d_L (source luminosity distance):{self.d_L.cgs:.2e}\n"
            + f" - delta_D (blob Doppler factor): {self.delta_D:.2e}\n"
            + f" - Gamma (blob Lorentz factor): {self.Gamma:.2e}\n"
            + f" - Beta (blob relativistic velocity): {self.Beta:.2e}\n"
            + f" - theta_s (jet viewing angle): {self.theta_s:.2e}\n"
            + f" - B (magnetic field tangled to the jet): {self.B:.2e}\n"
            + f" - xi (coefficient for 1st order Fermi acceleration) : {self.xi:.2e}\n"
            + str(self.n_e)
        )

    def set_delta_D(self, Gamma, theta_s):
        """set the viewing angle and the Lorentz factor of the outflow to
        obtain a specific Doppler factor

        Parameters
        ----------
        Gamma : float
            Lorentz factor of the relativistic outflow
        theta_s : :class:`~astropy.units.Quantity`
            viewing angle of the jet
        """
        mu_s = np.cos(theta_s.to("rad").value)
        Beta = np.sqrt(1 - 1 / np.power(Gamma, 2))
        delta_D = 1 / (Gamma * (1 - Beta * mu_s))

        self.theta_s = theta_s
        self.mu_s = mu_s
        self.Gamma = Gamma
        self.Beta = Beta
        self.delta_D = delta_D

    def N_e(self, gamma):
        r"""number of electrons as a function of the Lorentz factor, 
        :math:`N_e(\gamma') = V_b\,n_e(\gamma')`"""
        return self.V_b * self.n_e(gamma)

    @property
    def n_e_tot(self):
        r"""total electrons density

        .. math::
            n_{e,\,tot} = \int^{\gamma'_{\rm max}}_{\gamma'_{\rm min}} {\rm d}\gamma' n_e(\gamma')
        """
        return np.trapz(self.n_e(self.gamma), self.gamma)

    @property
    def N_e_tot(self):
        r"""total number of electrons

        .. math::
            N_{e,\,tot} = \int^{\gamma'_{\rm max}}_{\gamma'_{\rm min}} {\rm d}\gamma' N_e(\gamma')
        """
        return np.trapz(self.N_e(self.gamma), self.gamma)

    @property
    def u_e(self):
        r"""total energy density in non-thermal electrons

        .. math::
            u_{e} = m_e c^2\,\int^{\gamma'_{\rm max}}_{\gamma'_{\rm min}} {\rm d}\gamma' \gamma' n_e(\gamma')
        """
        return mec2 * np.trapz(self.gamma * self.n_e(self.gamma), self.gamma)

    @property
    def W_e(self):
        r"""total energy in non-thermal electrons

        .. math::
            W_{e} = m_e c^2\,\int^{\gamma'_{\rm max}}_{\gamma'_{\rm min}} {\rm d}\gamma' \gamma' N_e(\gamma')
        """
        return mec2 * np.trapz(self.gamma * self.N_e(self.gamma), self.gamma)

    @property
    def U_B(self):
        r"""energy density of magnetic field

        .. math::
            U_B = B^2 / (8 \pi)
        """
        U_B = np.power(self.B_cgs, 2) / (8 * np.pi)
        return U_B.to("erg cm-3")

    @property
    def k_eq(self):
        """equipartition parameter: ratio between totoal electron energy density  
        magnetic field energy density, Eq. 7.75 of [DermerMenon2009]_"""
        return (self.u_e / self.U_B).to_value("")

    @property
    def P_jet_e(self):
        r"""jet power in electrons

        .. math::
            P_{jet,\,e} = 2 \pi R_b^2 \beta \Gamma^2 c u_e
        """
        prefactor = (
            2 * np.pi * np.power(self.R_b, 2) * self.Beta * np.power(self.Gamma, 2) * c
        )
        return (prefactor * self.u_e).to("erg s-1")

    @property
    def P_jet_B(self):
        r"""jet power in magnetic field

        .. math::
            P_{jet,\,B} = 2 \pi R_b^2 beta \Gamma^2 c \frac{B^2}{8\pi}
        """
        prefactor = (
            2 * np.pi * np.power(self.R_b, 2) * self.Beta * np.power(self.Gamma, 2) * c
        )
        return (prefactor * self.U_B).to("erg s-1")

    @property
    def u_ph_synch(self):
        r"""energy density of the synchrotron photons energy losses are:

        .. math::        
            (\mathrm{d}E/\mathrm{d}t)_{\mathrm{synch}} = 4 / 3 \sigma_T c U_B \gamma^2 

        the radiation stays an average time of :math:`(3/4) (R_b/c)` (the factor of 3/4 cames from averaging over a sphere), 
        so an e- with Lorentz factor :math:`\gamma` produces:

        .. math::        
            0.75\,(\mathrm{d}E/\mathrm{d}t)_{\mathrm{synch}}\,(R_b/c)\,/\,V_b 

        of radiation. We need to integrate over the electron spectrum  (and multiply back by V_b)

        .. math::        
            0.75\,\int n_e(\gamma) (\mathrm{d}E/\mathrm{d}t)_{\mathrm{synch}}  R_b  \mathrm{d}\gamma
        
        so

        .. math::        
            u_{\mathrm{synch}} = \sigma_T  U_B  R_b  \int n_e(\gamma) \, \gamma^2 \mathrm{d}\gamma

        WARNING: this does not take into account SSA!
        """
        u_ph = (
            sigma_T.cgs
            * self.U_B
            * self.R_b
            * np.trapz(np.power(self.gamma, 2) * self.n_e(self.gamma), self.gamma)
        )
        return u_ph.to("erg cm-3")

    def plot_n_e(self, ax=None, gamma_power=0):
        """plot the  electron distribution
        
        Parameters 
        ----------
        ax : :class:`~matplotlib.axes.Axes`, optional
            Axis
        gamma_power : float
            power of gamma to raise the electron distribution
        """
        load_mpl_rc()
        ax = plt.gca() if ax is None else ax

        ax.loglog(self.gamma, np.power(self.gamma, gamma_power) * self.n_e(self.gamma))
        ax.set_xlabel(r"$\gamma$")
        if gamma_power == 0:
            ax.set_ylabel(r"$n_e(\gamma)\,/\,{\rm cm}^{-3}$")
        else:
            ax.set_ylabel(
                r"$\gamma^{"
                + str(gamma_power)
                + r"}$"
                + r"$\,n_e(\gamma)\,/\,{\rm cm}^{-3}$"
            )
        return ax
