from pylab import *

def B(T):
    return 5.67e-8*T**4

def B_eff_mu(B1, B2, tau, mu):
    dtau = tau/mu
    t = exp(-tau/mu)
    B_eff = (B1 * t - B2) / (t - 1.) + (B1 - B2) / dtau
    #B_eff = ((-B1 + (B2-B1) / tau * mu) * t + B1 + (B2-B1) / tau * (tau-mu)) / (1.-t)
    return B_eff * (1.-t)

def B_eff(B1, B2, tau, Nmu=100):
    '''
    Integrate[((-B1 + (B2 - B1)/tau*x)*Exp[tau/x] + B1 + (B2 - B1)/tau*(tau - x)), {x, 0.0000001, 1 - 0.0000001} ,
              Assumptions ->  B1 > 0 && B2 > 0 && tau > 0 && B1 is Real && B2 is Real && tau is Real]
    '''
    dmu = 1./Nmu
    B_eff = 0
    for mu in arange(dmu/2,1.,dmu):
        B_eff += B_eff_mu(B1/pi, B2/pi, tau, mu) * mu * dmu

    return B_eff * 2*pi

def schwarzschild_tau(T1, T2, tau=1., Nmu=100):
    Nlay = 1000
    dtau = tau/Nlay
    Tlev = linspace(T1,T2,Nlay+1)
    Blay = B( (Tlev[1:]+Tlev[:-1])/2 ) / pi
    dmu = 1./Nmu

    B_eff= 0
    for mu in arange(dmu/2,1.,dmu):
        L = 0 # lower bound
        t = exp(-dtau/mu) # transmission
        for k in range(Nlay):
            L = L * t + (1.-t) * Blay[k] # extinction plus emission with mean temperature
        B_eff += L * mu * dmu

    return B_eff * 2*pi

def B_schwarz(B1, B2, tau):
    return np.array([schwarzschild_tau(B1, B2, _) for _ in tau])

def B_2str(B1, B2, tau):
    Nmu = 50
    dmu = 1./Nmu
    t = np.mean([exp(-tau/mu) for mu in arange(dmu/2,1.,dmu)], axis=0)
    return ( B1 * t + B2*(1.-t) ) * (1.-t)

def B_Tmean(T1,T2, tau):
    Nmu = 50
    dmu = 1./Nmu
    t = np.mean([exp(-tau/mu) for mu in arange(dmu/2,1.,dmu)], axis=0)
    return B((T1+T2)/2) * (1.-t)


figure(1, figsize=(7,5))
clf()

tau = linspace(0,6,1e2)

T1, T2 = 288, 290
B1, B2, Bm = B(T1), B(T2), B((T1+T2)/2)

#axhline(B1, color='.5', linestyle='--')
axhline(B2, color='.5', linewidth=2, linestyle='--')
#annotate(r'$\rm B(T1)$', xy=(tau.max(), B1), xytext=(tau.max()*1.05, B1+abs(B1-B2)/20), arrowprops=dict(facecolor='black', shrink=0.01, width=.5, headwidth=4),)
annotate(r'$\rm B(T2)$', xy=(tau.max(), B2), xytext=(tau.max()*1.05, B2+abs(B1-B2)/20), arrowprops=dict(facecolor='black', shrink=0.01, width=.5, headwidth=4),)

#axhline(Bm, color='.5', linestyle='-', label=r'$\rm B(T_{mean}$)')
#annotate(r'$\rm B(T_{mean}$)', xy=(tau.max(), Bm), xytext=(tau.max()*1.05, Bm+abs(B1-B2)/20), arrowprops=dict(facecolor='black', shrink=0.01, width=.5, headwidth=4),)

plot(tau, B_Tmean(T1, T2, tau), color='.1', linestyle='-', linewidth=1, label=r'$\rm B_{eff}(T_{mean})$')
plot(tau, B_eff_mu(B1, B2, tau, 1.0), color='red', label=r'$\rm B_{eff}(\mu=1.0)$')
plot(tau, B_eff_mu(B1, B2, tau, 0.5), color='royalblue', label=r'$\rm B_{eff}(\mu=0.5)$')
plot(tau, B_eff_mu(B1, B2, tau, 0.6427), color='.2', linewidth=2, linestyle=':', label=r'$\rm B_{eff}(\mu=cos(50^\circ))$')

plot(tau, B_2str(B1, B2, tau), color='purple', linewidth=3, linestyle='--', label=r'$\rm B_{eff,2str} = B_1 * a_{11} + B_2*(1-a_{11})$')

plot(tau, B_schwarz(T1, T2, tau), color='green', linestyle='--', linewidth=5, label=r'$\rm schwarzschild\ numerical$')
plot(tau, B_eff(B1, B2, tau), color='orange', linewidth=1, label=r'$\int_{0}^{1} \rm B_{eff}(\mu)$')

ylim(0, B2*1.05)

legend(framealpha=.8,loc='best',fontsize='medium')
ylabel(r'$\rm eff.\ Planck\ emission\ [W/m^2]$')
xlabel(r'$\rm optical\ depth\ \tau$')

savefig('eff_emission.pdf', bbox_inches='tight')