import numpy as np
import astropy.constants as const
import astropy.units as u
from astropy.coordinates import Distance


# every variable indicated with capital letters is unitless
# will be used in SED computations for speed-up
E = const.e.gauss.value
H = const.h.cgs.value
C = const.c.cgs.value
ME = const.m_e.cgs.value
MEC = (const.m_e * const.c).cgs.value
MEC2 = (const.m_e * const.c * const.c).cgs.value
G = const.G.cgs.value
M_SUN = const.M_sun.cgs.value
K_B = const.k_B.cgs.value
SIGMA_SB = const.sigma_sb.cgs.value
EMISSIVITY_UNIT = "erg s-1"
SED_UNIT = "erg cm-2 s-1"


__all__ = ["SSDisk", "SphericalShellBLR", "RingDustTorus"]


def I_nu_bb(nu, T):
    """Black-Body intensity :math:`I_{\\nu}^{bb}`, Eq. 5.14 of [DermerMenon2009]_.
    Unitless parameters to speed-up calculations.

    Parameters
    ----------
    nu : :class:`~nump.ndarray`
        array of the frequencies (in Hz) to evaluate the Black Body intensity
    T : float or :class:`~nump.ndarray`
        temperature of the Black Body in K (might be function of the radius)
    """
    num = 2 * H * np.power(nu, 3)
    denum = np.power(C, 2) * (np.exp(H * nu / (K_B * T)) - 1)
    return num / denum


class SSDisk:
    """[Shakura1973]_ accretion disk.

    Parameters
    ----------
    M_BH : :class:`~astropy.units.Quantity`
        Black Hole mass    
    L_disk : :class:`~astropy.units.Quantity`
        luminosity of the disk 
    eta : float
        accretion efficiency
    R_in : :class:`~astropy.units.Quantity`
        inner disk radius
    R_out : :class:`~astropy.units.Quantity`
        outer disk radius
    """

    # properties with underscore are unitless, for fast computation

    def __init__(self, M_BH, L_disk, eta, R_in, R_out):
        self.type = "SSDisk"
        # masses and luminosities
        self.M_BH = M_BH.cgs
        self._M_BH = self.M_BH.value
        self.M_8 = self.M_BH.value / (1e8 * M_SUN)
        self.L_Edd = 1.26 * 1e46 * self.M_8 * u.Unit("erg s-1")
        self.L_disk = L_disk.cgs
        self._L_disk = self.L_disk.value
        # fraction of the Eddington luminosity at which the disk is accreting
        self.l_Edd = (self.L_disk / self.L_Edd).decompose()
        self.eta = eta
        self.m_dot = (self.L_disk / (self.eta * const.c * const.c)).cgs
        self._m_dot = self.m_dot.value
        # gravitational radius
        self._R_g = G * self.M_BH.value / np.power(C, 2)
        self.R_g = self._R_g * u.cm
        self.R_in = R_in.cgs
        self.R_out = R_out.cgs
        # for fast computation of the temperature dependency
        self._R_in = self.R_in.value
        self._R_out = self.R_out.value

    def __str__(self):
        summary = (
            f"* Shakura Sunyaev accretion disk:\n"
            + f" - M_BH (central black hole mass): {self.M_BH:.2e}\n"
            + f" - L_disk (disk luminosity): {self.L_disk:.2e}\n"
            + f" - eta (accretion efficiency): {self.eta:.2e}\n"
            + f" - dot(m) (mass accretion rate): {self.m_dot:.2e}\n"
            + f" - R_in (disk inner radius): {self.R_in:.2e}\n"
            + f" - R_out (disk inner radius): {self.R_out:.2e}"
        )
        return summary

    def _phi_disk(self, R):
        """Radial dependency of disk temperature
        Eq. 63 in [Dermer2009]_.

        Parameters
        ----------
        R : :class:`~numpy.ndarray`
            radial coordinate along the disk, in cm
            dimensionless to speed-up computation in the external Compton SED
        """
        return 1 - np.sqrt(self._R_in / R)

    def _phi_disk_mu(self, mu, r):
        """same as _phi_disk but computed with cosine of zenith mu and distance
        from the black hole r. Eq. 67 in [4]."""
        R = r * np.sqrt(np.power(mu, -2) - 1)
        return self._phi_disk(R)

    def _epsilon(self, R):
        """Monochromatic approximation for the mean photon energy at radius R
        of the accretion disk. Eq. 64 in [Dermer2009]_. R is unitless.
        """
        _term_1 = np.power(self.l_Edd / (self.M_8 * self.eta), 1 / 4)
        _term_2 = np.power(R / self._R_g, -3 / 4)
        _prefactor = 2.7 * 1e-4
        return _prefactor * _term_1 * _term_2

    def _epsilon_mu(self, mu, r):
        """same as _epsilon but computed with cosine of zenith mu and distance
        from the black hole r. Eq. 67 in [Dermer2009]_."""
        R = r * np.sqrt(np.power(mu, -2) - 1)
        return self._epsilon(R)

    def T(self, R):
        """Temperature of the disk at distance R. Eq. 64 in [Dermer2009]_."""
        value = MEC2 / (2.7 * K_B) * self._epsilon(R.value)
        return value * u.K

    def sed_flux(self, nu, z, mu_s):
        """Black Body SED generated by the SS Disk:

        .. math::
            \\nu F_{\\nu} \, [\mathrm{erg}\,\mathrm{cm}^{-2}\,\mathrm{s}^{-1}]
        
        ----------
        nu : :class:`~astropy.units.Quantity`
            array of frequencies, in Hz, to compute the sed, **note** these are 
            observed frequencies (observer frame).
        z : float
            redshift of the galaxy, to correct the observed frequencies and to 
            compute the flux once the distance is obtained
        mu_s : float
            cosine of the angle between the disk axis and the line of sight,
            same as the jet viewing angle
        """
        nu = nu.to("Hz").value * (1 + z)
        d_L = Distance(z=z).to("cm").value
        prefactor = 2 * np.pi * mu_s / np.power(d_L, 2)
        R = np.logspace(np.log10(self._R_in), np.log10(self._R_out)) * u.cm
        T = self.T(R)
        _R = R.value.reshape(R.size, 1)
        _T = T.value.reshape(T.size, 1)
        _nu = nu.reshape(1, nu.size)
        _integrand = _R * I_nu_bb(_nu, _T)
        return prefactor * nu * np.trapz(_integrand, R, axis=0) * u.Unit(SED_UNIT)


class SphericalShellBLR:
    """Spherical Shell Broad Line Region, from [Finke2016]_.
    Each line is emitted from an infinitesimally thin spherical shell. 

    Parameters
    ----------
    disk : :class:`~agnpy.targets.SSDisk`
        disk whose radiation is being reprocessed by the BLR
    xi_line : float
        fraction of the disk radiation reprocessed by the BLR
    epsilon_line : float
        dimensionless energy of the emitted line
    R_line : :class:`~astropy.units.Quantity`
        radius of the BLR spherical shell
    """

    def __init__(self, disk, xi_line, epsilon_line, R_line):
        self.type = "SphericalShellBLR"
        self.parent_disk = disk
        self.xi_line = xi_line
        self.epsilon_line = epsilon_line
        self.R_line = R_line.cgs
        self._R_line = self.R_line.value

    def __str__(self):
        summary = (
            f"* Spherical Shell Broad Line Region:\n"
            + f" - M_BH (central black hole mass): {self.parent_disk.M_BH:.2e}\n"
            + f" - L_disk (accretion disk luminosity): {self.parent_disk.L_disk:.2e}\n"
            + f" - xi_line (fraction of the disk radiation reprocessed by the BLR): {self.xi_line:.2e}\n"
            + f" - epsilon_line (dimensionless energy of the emitted line): {self.epsilon_line:.2e}\n"
            + f" - R_line (radius of the BLR shell): {self.R_line:.2e}\n"
        )
        return summary


class RingDustTorus:
    """Dust Torus as infinitesimally thin annulus, from [Finke2016]_.
    For the Compton scattering monochromatic enission at the peak energy of the 
    Black Body spectrum is considered.

    Parameters
    ----------
    disk : :class:`~agnpy.targets.SSDisk`
        disk whose radiation is being reprocessed by the Torus
    xi_dt : float
        fraction of the disk radiation reprocessed
    epsilon_dt : float
        dimensionless energy peak of the black body distribution
    R_dt : :class:`~astropy.units.Quantity`
        radius of the Torus, if not specified the saturation radius of Eq. 96 in
        [Finke2016]_ will be used
    """

    def __init__(self, disk, xi_dt, epsilon_dt, R_dt=None):
        self.type = "RingDustTorus"
        self.parent_disk = disk
        self.xi_dt = xi_dt
        self.epsilon_dt = epsilon_dt
        # dimensionless temperature of the Dust Torus
        self.Theta = self.epsilon_dt / 2.7
        # temperatue in K
        self.T_dt = self.Theta * ((const.m_e * const.c * const.c) / const.k_B).to("K")
        self._T_dt = self.T_dt.value
        # if the radius is not specified use saturation radius Eq. 96 of [Finke2016]_
        if R_dt is None:
            self.R_dt = (
                3.5
                * 1e18
                * np.sqrt(self.parent_disk._L_disk / 1e45)
                * np.power(self._T_dt / 1e3, -2.6)
            ) * u.cm
        else:
            self.R_dt = R_dt.cgs
        self._R_dt = self.R_dt.value

    def __str__(self):
        summary = (
            f"* SRing Dust Torus:\n"
            + f" - M_BH (central black hole mass): {self.parent_disk.M_BH:.2e}\n"
            + f" - L_disk (accretion disk luminosity): {self.parent_disk.L_disk:.2e}\n"
            + f" - xi_dt (fraction of the disk radiation reprocessed by the torus): {self.xi_dt:.2e}\n"
            + f" - epsilon_dt (dimensionless temperature of the dust torus): {self.epsilon_dt:.2e}\n"
            + f" - R_dt (radius of the torus): {self.R_dt:.2e}\n"
        )
        return summary

    def sed_flux(self, nu, z):
        """Black Body SED generated by the Dust Torus:

        .. math::
            \\nu F_{\\nu} \, [\mathrm{erg}\,\mathrm{cm}^{-2}\,\mathrm{s}^{-1}]
        
        ----------
        nu : :class:`~astropy.units.Quantity`
            array of frequencies, in Hz, to compute the sed, **note** these are 
            observed frequencies (observer frame).
        z : float
            redshift of the galaxy, to correct the observed frequencies and to 
            compute the flux once the distance is obtained
        """
        nu = nu.to("Hz").value * (1 + z)
        d_L = Distance(z=z).to("cm").value
        prefactor = np.pi * np.power(self._R_dt / d_L, 2)
        sed = prefactor * nu * I_nu_bb(nu, self._T_dt)
        return sed * u.Unit("erg cm-2 s-1")
