# -*- coding: utf-8 -*-
"""
1) save this script in the same folder as your data (spectra in txt files from Notepd)

2) set/change the filenames here:
"""
file_back = '20221223_background.txt'   # filename background
file_data = '20221223_heat_460C.txt'   # filename of data
"""
Note: you can measure the background (spectrum when it is completely dark) just
once, and use reuse it for each measurement.

3) set de minimum and maximum value for the x-axis:
"""
min_wave = 200e-9   # 900 nm
max_wave = 1700e-9  # 1700 nm
"""
4) set the reference temperature lines
"""
T_min = 770
T_max = 1400
num_lines = 11
"""
5) set the scale of the theorethical curves
smaller values make the curves more steep at the start
"""
scale = 30e+4
"""
Using these reference temperature lines and by carefully adjusting the scale
and the min and max temperatures you can determine the temperature of
the object you measured the Black Body Radiation from. The curves of the theory
lines need to fit the steep slope of the onset on the IR spectrum as best as
possible.

At what temperature does the spectrum fit best?

5) Run the script (green play button at the top). Change the file-names and
values above to get to your answer!

6) Plot data from visible spectrum; switch between True and False
"""
plot_vis = False



"""
#####################################################################
Below is the script; do NOT make any changes here unless
you know what you are doing!
#####################################################################
"""
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.pyplot import cm
import scipy.constants as const
import os

# set folder to folder of script
file = os.path.dirname(__file__)
os.chdir(file)

k = const.k   # 1.3806e-23 Boltzmann constant
h = const.h     # 6.6261e-34 Planck constant
c = const.c     # 2.9979e+08 speed of light
pi = np.pi      # 3.1415e+00 pi

def PL(wavelength, T):
    """
    Planck's Law (quantum)

    Parameters
    ----------
    wavelength : int, float, np.array
        wavelength in m (or nm?)
    T : int, float
        temperature in K

    Returns
    -------
    float, np.array
        spectral radiance in W·sr−1·m−3
        (or spectral radiance in W·sr−1·m−2·nm−1?)
    """
    return (2*pi*h*c**2)/((wavelength**5)*np.exp(((h*c)/(wavelength*k*T))-1))

# load background
filename = file_back

f = open(filename, 'r')
data_list = f.readlines()            # reads each line into a new element of a list

data1 = []
for line in data_list:
    data1.append(line.split())       # splits data in a str into elements at every space

data_b_vis_x = []
data_b_vis_y = []
data_b_IR_x = []
data_b_IR_y = []
for line in data1:
    data_b_vis_x.append(float(line[0]))
    data_b_vis_y.append(float(line[1]))
    if len(line) == 4:
        data_b_IR_x.append(float(line[2]))
        data_b_IR_y.append(float(line[3]))

data_b_vis = np.array((data_b_vis_x,data_b_vis_y))
data_b_IR = np.array((data_b_IR_x,data_b_IR_y))
data_b_vis = data_b_vis.T
data_b_IR = data_b_IR.T

# load data
filename = file_data

f = open(filename, 'r')
data_list = f.readlines()            # reads each line into a new element of a list

data1 = []
for line in data_list:
    data1.append(line.split())       # splits data in a str into elements at every space

data_vis_x = []
data_vis_y = []
data_IR_x = []
data_IR_y = []
for line in data1:
    data_vis_x.append(float(line[0]))
    data_vis_y.append(float(line[1]))
    if len(line) == 4:
        data_IR_x.append(float(line[2]))
        data_IR_y.append(float(line[3]))

data_vis = np.array((data_vis_x,data_vis_y))
data_IR = np.array((data_IR_x,data_IR_y))
data_vis = data_vis.T
data_IR = data_IR.T

# subtract background
data_vis[:,1] = data_vis[:,1] - data_b_vis[:,1]
data_IR[:,1] = data_IR[:,1] - data_b_IR[:,1]

# align intensity spectra
factor_vis = np.average(data_vis[data_vis[:,0] > np.min(data_IR[:,0]),1])
factor_IR = np.average(data_IR[data_IR[:,0] < np.max(data_vis[:,0]),1])
factor = factor_vis/factor_IR

data_IR[:,1] = data_IR[:,1]*factor      # correct IR data with factor

# plot with temperature curves from theory
T_range = np.linspace(T_min,T_max,num=num_lines)
x = np.linspace(min_wave,max_wave,num=101)

colors = plt.cm.nipy_spectral(np.linspace(0,1,num_lines))

plt.figure(figsize=[12,6])
plt.plot(data_IR[:,0]*10E-10,data_IR[:,1],'.k',label='data')
if plot_vis == True:
    plt.plot(data_vis[:,0]*10E-10,data_vis[:,1],'.k')

for i, temp in enumerate(T_range):
    plt.plot(x,PL(x,temp)/scale,label=int(temp), c=colors[i])
plt.ylim(np.min(data_vis[:,1]),1.1*np.max(data_IR[:,1]))
plt.xlim((min_wave,max_wave))
plt.xlabel('golflengte $λ$ in m')
plt.ylabel('spectrale stralings intensiteit $B_λ$ in W·sr$^{−1}$·m$^{−3}$')
plt.title('Quantum mechanics theorie (kleur) spectral radiance per wavelength')
plt.legend(title='temperatuur in K')
plt.savefig('theory.png',dpi=300)
plt.show()