#The following chunk of code imports the various
#libraries we need (numpy, matplotlib etc.)
from picamera import PiCamera
from time import sleep
from picamera.array import PiRGBArray
from matplotlib import pyplot as PLT
from matplotlib import patches as PAT
import numpy as np

#The following chunk of code sets up the PiCamera to give a preview
#of what the camera can 'see'
camera = PiCamera()
camera.start_preview()
sleep(5) #Change the number here in brackets for a longer/shorter preview
camera.stop_preview()

#The following chunk of code captures the image into a 3-dimensional array
#of RGB values. The array (called output) is the RGB values of each pixel in the camera sensor
#in the camera sensor. The RGB values are summed for each pixel to give the
#2-dimensioanl array (summed). This is ment to represent 'intensity' of the 
#light falling on each pixel rather than colour.
output= PiRGBArray(camera)
camera.capture(output, 'rgb')
Summed = np.flipud(np.sum(output.array,axis=2))

#The following 3 lines of code set up where on the PiCamera image the spectrum
#will be taken from. Feel free to change these values to suit your or the
#students needs
Spectrum_Line_Height = 540 #This contorols the height of the white line
Spectrum_Line_Minimum_Horizontal_Pixel = 500 #This controls the start pixel of the spectrum
Spectrum_Line_Maximum_Horizontal_Pixel = 1100#This controls the end pixel of the spectrum

#THe following two chunks of code plot the image from the PiCamera, the white
#line (which is indicative of where the spectrum is taken from)). and the
#2-dimensional plot showing the spectrum of the light source
PLT.subplot(1,2,1)
PLT.imshow(np.rot90(output.array,2))
PLT.plot([0, 1920], [Spectrum_Line_Height, Spectrum_Line_Height],'w-')
PLT.xlim((output.array.shape[1],0))
PLT.ylim((0,output.array.shape[0]))
PLT.title('PiCamera Image from Spectrometer')

PLT.subplot(1,2,2)
PLT.plot(Summed[Spectrum_Line_Height,:])
PLT.xlim((Spectrum_Line_Minimum_Horizontal_Pixel,Spectrum_Line_Maximum_Horizontal_Pixel))
PLT.title('Spectrum from Spectrometer Image')
PLT.xlabel('Pixel Number')
PLT.ylabel('Intensity')
PLT.show()