#!/usr/bin/python
# -*- coding: utf8 -*-

import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
from math import *

code_website = 'http://commons.wikimedia.org/wiki/User:Geek3/mplwp'
try:
    import mplwp
except ImportError, er:
    print 'ImportError:', er
    print 'You need to download mplwp.py from', code_website
    exit(1)

name = 'Mplwp_blackbody_nu_planck-wien-rj_5800K.svg'
fig = mplwp.fig_standard(mpl)
xlim = 0, 1400; fig.gca().set_xlim(xlim)
ylim = 0, 40; fig.gca().set_ylim(ylim)
#fig.gca().xaxis.set_major_locator(mpl.ticker.MultipleLocator(400))
mplwp.set_bordersize(fig, 56.5, 30.5, 24.5, 48.5)

h = 6.62607e-34
c = 2.99792e8
kB = 1.38065e-23
T = 5800

def planck(nu, T):
    f = nu * 1e12
    if f == 0.0 or h * f / (kB * T) > 1e2:
        return 0.0
    return 2 * h * f**3 / c**2 / (exp(h * f / (kB * T)) - 1.0)

def wien(nu, T):
    f = nu * 1e12
    if h * f / (kB * T) > 1e2:
        return 0.0
    return 2 * h * f**3 / c**2 * exp(-h * f / (kB * T))

def rayleigh_jeans(nu, T):
    f = nu * 1e12
    return 2 * f**2 * kB * T / c**2

x = np.linspace(xlim[0], xlim[1], 8001)
y_planck = [planck(xx, T) * 1e9 for xx in x]
y_wien = [wien(xx, T) * 1e9 for xx in x]
y_rj = [rayleigh_jeans(xx, T) * 1e9 for xx in x]
pp, = plt.plot(x, y_planck, '-', label='Planck', zorder=-1)
pw, = plt.plot(x, y_wien, '--', dashes=[10, 4], label='Wien', zorder=-2)
pr, = plt.plot(x, y_rj, ':', dashes=[3, 6], label='Rayleigh-Jeans', zorder=-3)

plt.xlabel(ur'$\nu$ [THz]')
plt.ylabel(ur'$B_{\nu}\;[kW\,sr^{-1}\,m^{-2}\,THz^{-1}]$')
plt.title('T = 5800 K')
plt.legend([pr, pp, pw], ['Rayleigh-Jeans', 'Planck', 'Wien'],
           loc='upper right', borderaxespad=1)
plt.savefig(name)
mplwp.postprocess(name)