import matplotlib.pyplot as plt
import numpy as np
import random
import pandas as pd 
%matplotlib inline

# define common constants
from scipy.constants import  Planck, speed_of_light, Boltzmann
h = Planck  # Js
c = speed_of_light  # m/s
kb = Boltzmann  # J/molec K

# use "lamda" to avoid conflict with python reserved word "lambda"
def rho(wavelength, T):
    """Calculate the emission from wavelength and temperature.

    Arguments::
    lamda: wavelength cm
    T: temperature K"""

    return 8 * np.pi * h * c / ( wavelength**5 * (np.exp(h*c/(wavelength*kb*T)) - 1))

wavelength =np.linspace(1*10**-7, 3*10**-6,1000)
wavelength

TK = 4000   # temperature of a K-type star
TG = 5500   # temperature of a G-type star
TF = 7000   # temperature of a F-type star
IntensityK = rho(wavelength,TK)
IntensityG = rho(wavelength,TG)
IntensityF = rho(wavelength,TF)

plt.figure()
plt.plot(wavelength,IntensityK,'g-')
plt.plot(wavelength, IntensityG,'k-')
plt.plot(wavelength, IntensityF, 'r-')

def rhonu(nu, T):
    """Calculates the radiance from frequency and temperature.

    Arguments
    nu: frequency Hz
    T: temperature K
    """

    return  (8 * np.pi * h * nu**3 / c**3) * (1 / (np.exp(h*nu/(kb*T)) - 1))

nu1 = c/200e-9
nu2 = c/900e-8     # frequency


x = np.linspace(nu1,nu2,1000)
y=rhonu(x,4000)

plt.figure()
plt.plot(x,rhonu(x,4000),'g')
plt.plot(x,rhonu(x,5000),'r')
plt.plot(x,rhonu(x,7000),'k')      # Intensity versus frequency


#rhonu(nu,4000)

wavelength =np.linspace(100e-9, 2000e-9, 1000)

plt.figure()
plt.plot(wavelength,rho(wavelength, 4000),'g')   # The function for intensity versus wavelength for a K-type star
plt.plot(wavelength,rho(wavelength, 5000),'r')   # The function for intensity versus wavelenght for a G-type star
plt.plot(wavelength,rho(wavelength,7000),'k')    # The function for intensity versus wavelength for a F-type star

from scipy.constants import Planck, speed_of_light, Boltzmann
h =  Planck  #Js
c = speed_of_light   #m.s
kb = Boltzmann    # J/molec K

def rhonu(nu, T):
    """
    Return emission intensity as function of frequency and temperature.
    
    Arguments::
    nu =  frequency Hz
    T =  temperature K
    """
    
    return 2*h*nu**2/c**2 * 1/(np.exp(h*nu/(kb*T)-1))

nus = np.linspace(1e10, 1e15, 100)
for T in [4000, 5000, 7000]:
    plt.plot(nus,rhonu(nus, T), label='%s K'%T)
plt.xlabel('Frequency/Hz')
plt.ylabel('Intensity')
plt.legend()

#wl = np.linspace(1e-8, 25e-7, 1000)
nus = np.linspace(0, 2e15, 100)   #Hz
for t in [7000, 5000, 4000]:
    y = rhonu(nus,t)
    plt.plot(nus, rhonu(nus, t), label=(f'{t}K'))
    mx = np.argmax(y)
    print(mx,nus[mx],y[mx])     # This operation will print all the lambda maximum frequency values and all the maximum y values for each varied temperature curve. 
    plt.scatter(nus[mx], y[mx], c='k')
plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))
plt.ticklabel_format(style='sci', axis='y', scilimits=(0,0))
plt.xlabel(r'frequency / $s^{-1}$')
plt.ylabel('Intensity')
plt.legend()

fn = 'data/ukg5iii.csv'
G = pd.read_csv(fn, skiprows=3, names=["wavelength", "intensity"])

# convert Angstroms to meters
wavelength = G.wavelength / 1e10 

# normalize intensity
em_star = G.intensity / np.max(G.intensity)

plt.plot(wavelength, em_star, 'k', lw=0.5)
plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))
plt.xlabel("Wavelength / m");
plt.ylabel("Intensity")
plt.title("Experimental Analysis of Intensity Plotted against Wavelength /m")

# This is the experimental analysis of the plot of Intensity versus wavelength from the imported data

plt.plot(wavelength, em_star)

T = 5300 # temperature of a K-type star

em_calc = rho(wavelength, T)

em_norm = em_calc/np.max(em_calc)

plt.plot(wavelength, em_norm)

plt.xlabel("wavelength /m")
plt.ylabel("Intensity")
plt.title("Calculated Results for Intensity Plotted against Wavelength /m")

# I made a normalization curves, while varying the temperature to see if I could get my calculated curve as close as possible to the experimental analysis curve. Then I overlayed my generated curve on the experimental analysis curve.

from scipy.optimize import curve_fit

def model(wavelength, T):
    em_calc =  rho(wavelength, T)  # calculated emission
    print(em_calc)
    em_norm = em_calc/np.max(em_calc)   # normalize calculated emission
    return em_norm

T = 4000  # initial estimation for temperature of a K-type star
popt, pcov = curve_fit(model, wavelength, T)    # for curve_fit, model comes first, then the x and y data of which we are comparing the model to.
popt

plt.plot(wavelength, em_star)
plt.plot(wavelength, model(wavelength, *popt))
plt.ticklabel_format(style='sci', axis='x', scilimits=(0,0))
plt.ylabel("Normalizaed Intensity")
plt.xlabel("Wavelength/m")

def wien(T):
    "lambda_max for Wein law; T is temperature in K."
    l_max = 2.89e-3 / T
    return l_max

temps = np.arange(1000,10001,1000)

planck_list =[]
wien_list = []
for temp in temps:
    y = rho(wavelength, temp)
    max_y = np.argmax(g)
    wavelength = wavelength[max_y]
    planck_list.append(wavelength)
    
    planck_list.append(wavelength[np.argmax(rho(wavelength,temp))])
    wein_list.append(wien(temp))
    
    print("Wien)\t\t\tPlanck")
    for w,p in zip(wien_list, planck_list):
        print("f{w:.3e} \t\t {p:.3e}")

plt.plot(wien_list, planck_list, 'o')
plt.xlabel("Wien max")
plt.ylabel("Planck max")

np.polyfit(wien_list,planck_list, 1)

import pandas as pd
G = pd.read_csv('data/ukg5iii.csv', skiprows=3, names=['wavelength','intensity'])
plt.plot(G.wavelength, G.intensity, 'k', lw=0.5)
plt.ylabel('Intensity')
plt.xlabel(r'Wavelength / $\AA$')