
import numpy as np
import astropy.units as u
import astropy.constants as con
from astropy.constants import c, k_B, h
from astropy.cosmology import FlatLambdaCDM
cosmo = FlatLambdaCDM(H0=70.0, Om0=0.30) 

c=c.value
k_B = k_B.value
h = h.value


# DEF_BETA = 1.8 #default beta
# DEF_ALPHA = 2.0 #default alpha
# DEF_T = 50
Tcmb0 = 2.75

# note in all cases, lambda (l) is in microns and rest frame. These functions return model values in Janskys 

def _ot_mbb(l, Nbb, beta, T, z,l0=200):
    """Optically thin modified blackbody function"""
    Tcmbz = Tcmb0*(1+z)
    Tz = (T**(4+beta)  +  (Tcmb0)**(4+beta) * ((1+z)**(4+beta)-1) ) **(1/(4+beta))
    result = (1 - (planckbb(l,Tcmbz)/planckbb(l,Tz))) * (Tz/T)**(4+beta) \
            * 10.0**Nbb *(l0*1e-6/c)**beta * 2*h / c**2 * (c/(l*1.e-6))**(beta+3)/(np.exp(h*c/(l*1e-6*k_B*T))-1)
    return result

def _go_mbb(l, Nbb, beta, T, z,l0=200):
    """General opacity modified blackbody function"""
    Tcmbz = Tcmb0*(1+z)
    Tz = (T**(4+beta)  +  (Tcmb0)**(4+beta) * ((1+z)**(4+beta)-1) ) **(1/(4+beta))
    ### general opacity version
    result = (1 - (planckbb(l,Tcmbz)/planckbb(l,Tz))) * (Tz/T)**(4+beta) \
                * 10.0**Nbb * 2*h / c**2 * ((1.0 - np.exp(-(l0/l)**beta))*(c/(l*1.e-6))**3.0)/(np.exp(h*c/(l*1.e-6*k_B*T))-1.0)
    return result

def _powerlaw(l, Npl, alpha, l_c):
    """MIR power-law. Note that here, unlike Casey+2012, the normalization Npl does not implicitly include the factor l_c^alpha. 
    Instead, it is just the value of the MBB at l_c."""
    return Npl*(l/l_c)**alpha

def _approx_l_c(alpha, T, opthin=False, scale=1.0):
    """Approximate turnover wavelength as defined in Casey 2012"""
    if opthin: return scale*(T*(alpha*7.243e-5 + 1.905e-4))**-1.0
    return scale*(T*(alpha*7.243e-5 + 1.905e-4)+ (alpha*6.246 + 26.68)**-2.0)**-1.0

def _exact_l_c(N, beta, T, z, alpha, l0, opthin=False, scale=1.0):
    """Direct turnover wavelength as measured from mbb gradient"""
    l=np.logspace(1,3,1001)
    if opthin: y = _ot_mbb(l, N,beta,T,z,l0=l0)
    else: y = _go_mbb(l, N,beta,T,z,l0=l0)
    dy = np.log10(y[1:]/y[:-1])
    dx = np.log10(l[1:]/l[:-1])
    slope=dy/dx
    i_join = np.nanargmin(np.abs(slope-alpha)) # match alpha to slope
    l_c = scale*l[i_join]
    return l_c


def mbb_func(l, N=12,beta=1.8,T=35,z=0,alpha=2.0, l0=200, opthin=True, pl=False, pl_piecewise=False):
    """MBB function with optional powerlaw and variable opacity assumptions

    Args:
        l (float): the rest-frame wavelengths, in microns, at which to evaluate the model.
        N (float): log10 normalization factor of blackbody
        beta (float): emissivity index
        T (float): effective dust temperature
        z (float): the redshift of the model
        alpha (float): power-law slope
        l0 (float): turnover wavelength at which dust is optically thin
        opthin (bool): Whether or not the model should assume optically thin dust emission.
        pl (bool): Whether or not the model should include a MIR power law 
        pl_piecewise (bool): if the powerlaw should be attached piecewise (as in Drew&Casey 2022) or fall off exponentially (as in Casey+ 2012)

    Returns:
        float: the value(s) of the model in Jy at wavelengths ``l``, in microns
    """
    l = np.array(l,)
    if pl:
        # todo: add flexibility so piecewise is independent of scale
        scale=0.75
        if pl_piecewise: scale=1.0

        try:
            l_c = _exact_l_c(N, beta, T, z, alpha, l0=l0, opthin=opthin, scale=scale)
        except ValueError: l_c = _approx_l_c(alpha, T, opthin=opthin, scale=scale)
        # l_c = _approx_l_c(alpha, T, opthin=opthin, scale=scale)
        if opthin: 
            norm = np.atleast_1d(_ot_mbb(l_c, N,beta,T,z,l0=l0))
            mbb_y = np.atleast_1d(_ot_mbb(l, N,beta,T,z,l0=l0))
            pl_y = np.atleast_1d(_powerlaw(l, norm, alpha, l_c=l_c))
        else: 
            norm = np.atleast_1d(_go_mbb(l_c, N,beta,T,z,l0=l0))
            mbb_y = np.atleast_1d(_go_mbb(l, N,beta,T,z,l0=l0))
            pl_y = np.atleast_1d(_powerlaw(l, norm, alpha,l_c=l_c))
        if pl_piecewise: 
            mbb_y[l<l_c] = 0
            pl_y[l>=l_c] = 0
        else:
            pl_y *= np.exp(-(l/l_c)**2) #smooth fall off above l_c if not piecewise
        return np.squeeze(mbb_y + pl_y)
    else:
        if opthin: return _ot_mbb(l, N,beta,T,z,l0=l0) 
        else: return _go_mbb(l, N,beta,T,z,l0=l0) 


def planckbb(l,T): # note although this requires wavlength in micron, it represents B_nu
    return 2*h / c**2 * (c/(l*1e-6))**3 / (np.exp(h*c/(l*1e-6*k_B*T))-1)
