import matplotlib.pyplot as plt
import astropy.io.fits as fits
import astropy.modeling.blackbody as bb

#1A

f1 = fits.open('ckp00_3500.fits') #These lines are used to load in the .fits files.
f2 = fits.open('ckp00_5750.fits')
f3 = fits.open('ckp00_7500.fits')

data1 = f1[1].data #Here the necessary data is extracted.
data2 = f2[1].data
data3 = f3[1].data

wavelength1 = data1.field('wavelength') #The array with measured wavelengths is requested for each of the stars.
wavelength2 = data2.field('wavelength')
wavelength3 = data3.field('wavelength')

flux1 = data1.field('g00') #The array with measured flux is reuqested for each of the stars.
flux2 = data2.field('g00')
flux3 = data3.field('g05') #I use g05 instead of g00 here, because the g00 measurement isn't a valid spectrum.

plt.plot(wavelength1,flux1) #Plot the model spectrum of the star.
plt.xlim([3000,10000]) #We limit the wavelength to include only the optical part of the EM-spectrum.
plt.title('Model spectrum of star with T = 3500')
plt.xlabel('Wavelength (Ångstrom)')
plt.ylabel('Flux (erg/s/cm²/Å)')
plt.show()
plt.plot(wavelength2,flux2)
plt.xlim([3000,10000])
plt.title('Model spectrum of star with T = 5750')
plt.xlabel('Wavelength (Ångstrom)')
plt.ylabel('Flux (erg/s/cm²/Å)')
plt.show()
plt.plot(wavelength3,flux3)
plt.xlim([3000,10000])
plt.title('Model spectrum of star with T = 7500')
plt.xlabel('Wavelength (Ångstrom)')
plt.ylabel('Flux (erg/s/cm²/Å)')
plt.show()

#2A

plt.plot(wavelength1,flux1,label='Model spectrum')
plt.plot(wavelength1,bb.blackbody_lambda(wavelength1,3500),label='Blackbody curve')
#In the line above, the imported blackbody modeling is used to plot a blackbody curve, which is constructed using
#the array with wavelengths and the blackbody temperature of the star.
plt.vlines(8279,1e-16,1e9,label='Peak (Wiens law)',linestyle='--') #The calculated location of the peak according to Wien.
plt.xlim([1000,100000])
plt.ylim([1e-16,1e9])
plt.xscale('log') #Logaritmic scaling of the axis.
plt.yscale('log')
plt.legend(loc='lower right') #Legend to label the different plotted lines.
plt.title('Model spectrum of star with T = 3500 K, with overplotted blackbody curve')
plt.xlabel('Wavelength (Ångstrom)')
plt.ylabel('Flux (erg/s/cm²/Å)')
plt.show()

plt.plot(wavelength2,flux2,label='Model spectrum')
plt.plot(wavelength2,bb.blackbody_lambda(wavelength2,5750),label='Blackbody curve')
plt.vlines(5039.6,1e-16,1e9,label='Peak (Wiens law)',linestyle='--')
plt.xlim([1000,100000])
plt.ylim([1e-16,1e9])
plt.xscale('log')
plt.yscale('log')
plt.legend(loc='lower right')
plt.title('Model spectrum of star with T = 5750 K, with overplotted blackbody curve')
plt.xlabel('Wavelength (Ångstrom)')
plt.ylabel('Flux (erg/s/cm²/Å)')
plt.show()

plt.plot(wavelength3,flux3,label='Model spectrum')
plt.plot(wavelength3,bb.blackbody_lambda(wavelength3,7500),label='Blackbody curve')
plt.vlines(3864,1e-16,1e9,label='Peak (Wiens law)',linestyle='--')
plt.xlim([1000,100000])
plt.ylim([1e-16,1e9])
plt.xscale('log')
plt.yscale('log')
plt.legend(loc='lower right')
plt.title('Model spectrum of star with T = 7500 K, with overplotted blackbody curve')
plt.xlabel('Wavelength (Ångstrom)')
plt.ylabel('Flux (erg/s/cm²/Å)')
plt.show()