#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Mon Jul  6 16:00:24 2020

@author: sylviaploeckinger
"""

import numpy as np
import matplotlib.pyplot as plt
cmap = plt.cm.get_cmap('coolwarm')

spectral_class = ["O", "B", "A", "F", "G", "K", "M"]
temp_max       = np.asarray([6e4, 3e4, 1e4, 7.5e3, 6.2e3, 5.3e3, 3.9e3])
temp_min       = np.zeros_like(temp_max)
temp_min[0:-1] = temp_max[1:]
temp_min[-1]   = 2.7e3
temp_mean      = 0.5 *(temp_max + temp_min)
color = cmap(np.linspace(0,1,len(spectral_class)))

# constants
h  =  4.135667696e-15	    # eV s
kB =  8.617333262145e-5   # eV K−1
c  =  2.998e10            # cm s-1
nm_to_cm = 1.e-7          # 1 nm = <nm_to_cm> cm
eV_to_erg = 1.60218e-12   # 1 eV = <eV_to_erg> erg
W_to_erg_per_s = 1.e7     # 1 W = <W_to_erg_per_s> erg/s

def blackbody(lambda_nm, T):   # W cm-3 sr-1
    return eV_to_erg/W_to_erg_per_s * (2. * h * c * c / np.power(lambda_nm * nm_to_cm,5))/\
           (np.exp(h * c / (lambda_nm * nm_to_cm * kB * T)) - 1.)
           
           
lambda_nm_min = 10.
lambda_nm_max = 1000.
nbins    = 1000

lambda_nm = np.linspace(lambda_nm_min, lambda_nm_max, nbins) 


fig = plt.figure()
ax = plt.subplot(111)
ax.set_ylim(-4,15)
ax.set_yticks(np.arange(-4, 16, 2))
ax.set_title ('Blackbody radiation')
ax.set_xlabel('Wavelength [nm]')
ax.set_ylabel('log Spectral radiance [W cm$^{-3}$ sr$^{-1}$]')

for i in range (len(temp_max)):
    B = np.log10(blackbody(lambda_nm, temp_mean[i]))
    indx = np.where(B == B.max())[0]
    ax.axvline(lambda_nm[indx], color = color[i], ymin = 0.9)
    
for i in range (len(temp_max)):
    B = np.log10(blackbody(lambda_nm, temp_mean[i]))
    ax.plot(lambda_nm, B, label = "%s (%.0f K)"%(spectral_class[i], temp_mean[i]), color = color[i])
    
ax.legend()

fig.savefig('blackbody.png', dpi = 150)