#Tutorial 3 Solutions 
import numpy as np
import matplotlib.pyplot as plt 

#Load up neon spectrum
spectrum = np.loadtxt('neon.txt')
spectrum = np.transpose(spectrum)
pix = spectrum[0]
signal = spectrum[1]

#Find the peaks 

threshhold = 18000 #You can just pick slightly lower than the lowest peak you want to centroid
peaks = [] #x positions of the peaks, or rather, their index
for i in range(len(signal)-1):
    if (signal[i] > signal[i+1]) and (signal[i]>signal[i-1]) and (signal[i]>threshhold):
        if (signal[i] > signal[i-2]) and (signal[i] > signal[i+2]):
            peaks.append(i)
'''
The above for loop iterates through all values up to the second to last (so that i+2 actually exists). I have the if statements on two lines
just for ease of reading, it couldve been one long if statement too. The second line just introduces more robust peak finding (peak has to be
greater than the 2 spots on either side, not just one. I append to peaks the index of those peaks. I choose a threshold of 18000 which catches 
most peaks, feel free to experiment.
'''
#calculate the centroid around each peak 
window = 5 #[pixels]
centroids = [] #Values for all the centroids
for i in peaks:
    x_range = pix[i-window:i+window]
    I_range = signal[i-window:i+window]
    x_range = np.array(x_range)
    I_range = np.array(I_range)
    xcm = np.sum((x_range*I_range)) / np.sum(I_range)
    centroids.append(xcm) 

'''
In this block, I initialize an empty list to hold the centroids, then for each index of a peak, I construct a range over which to centroid
(by using my window to go 5 elements in either direction), and then use the xcm formula from the handout. I convert x_range and I_range to 
arrays because I multiply them together in the step below which is an array operation.
'''
def plot_vert(x): 
    '''
    Just plots vertical lines, in blue dashes
    '''
    plt.axvline(x, color='blue', ls='-.')
    
for i in centroids: #Call my plotting function on every centroid
    plot_vert(i)
    
    
plt.plot(pix, signal, 'r', label='Spectrum') #Plot the actual spectrum

plt.show() #Show it