# Code to fit Gaussian curves to multiple peaks on the same spectrum - Will Burrows - 20161027

import numpy as np
import matplotlib.pyplot as plt
import pyfits
from scipy.ndimage import median_filter
from scipy.optimize import curve_fit

# standard image processing first steps
f = pyfits.open("Calib_Cd_1.50.fit")
d = f[0].data
h = f[0].header
gain = h.get("EGAIN")
px = d.sum(0)
py = d.sum(1)
gx = px * gain
gy = py * gain
filtered = median_filter(d, size=4)
fpx = filtered.sum(0)
fpy = filtered.sum(1)
fgx = fpx * gain
fgy = fpy * gain

pixels = np.arange(0, len(fgx), 1)

fig = plt.figure(1)
plt.subplot(2, 1, 1)
plt.plot(pixels, fgx)

plt.subplot(2, 1, 2)
plt.plot(pixels, fgx, 'b')

def gauss(x, a, x0, sigma, bg):
    return a * np.exp(-(x - x0)**2 / (2 * sigma**2)) + bg
    
#n = len(pixels)
#mu = sum(pixels * fgx) / n
#sigma = sum(fgx * (pixels - mu)**2 ) / n

# takes array of arrays of form [peak 1 features, peak 2 features, etc...]
# peak features array [amplitude, pixel pos, width, background]
def gauss_fits(peaks_array) :  
    for i in range(0, len(peaks_array)):
        
        #popt, pcov = curve_fit(gauss, pixels, fgx, [peaks_array[i][0], peaks_array[i][1], peaks_array[i][2], peaks_array[i][3]])
        #plt.plot(pixels, gauss(pixels, popt[0], popt[1], popt[2], popt[3]), 'r')
        
        #temp_xrange = np.arange(peaks_array[i][1] - 50, peaks_array[i][1] + 50, 1)
        #temp_yrange = fgx[peaks_array[i][1] - 50 : peaks_array[i][1] + 50]
        #popt, pcov = curve_fit(gauss, temp_xrange, temp_yrange, [peaks_array[i][0], peaks_array[i][1], peaks_array[i][2], peaks_array[i][3]])
        #plt.plot(temp_xrange, gauss(temp_xrange, popt[0], popt[1], popt[2], popt[3]), 'r')
        
        temp_xrange = []
        temp_yrange = []
        if peaks_array[i][1] < 50:
            temp_xrange = np.arange(0, peaks_array[i][1] + 50, 1)
        elif peaks_array[i][1] > len(pixels) - 50:
            temp_xrange = np.arange(peaks_array[i][1] - 50, len(pixels))
        elif peaks_array[i][1] >= 50:
            temp_xrange = np.arange(peaks_array[i][1] - 50, peaks_array[i][1] + 50, 1)
        temp_yrange = fgx[temp_xrange]
        popt, pcov = curve_fit(gauss, temp_xrange, temp_yrange, [peaks_array[i][0], peaks_array[i][1], peaks_array[i][2], peaks_array[i][3]])
        plt.plot(temp_xrange, gauss(temp_xrange, popt[0], popt[1], popt[2], popt[3]), 'r')
        
        #temp_xrange = []
        #temp_yrange = []
        #for j in range(1, len(peaks_array) - 1):
        #    next_peak_dif = peaks_array[j + 1][1] - peaks_array[j][1]
        #    previous_peak_dif = peaks_array[j][1] - peaks_array[j - 1][1]
        #
        #if next_peak_dif >= 50 & previous_peak_dif >= 50:
        #    if peaks_array[i][1] < 50:
        #        temp_xrange = np.arange(0, peaks_array[i][1] + 50, 1)
        #    elif peaks_array[i][1] > len(pixels) - 50:
        #        temp_xrange = np.arange(peaks_array[i][1] - 50, len(pixels))
        #    elif peaks_array[i][1] >= 50:
        #        temp_xrange = np.arange(peaks_array[i][1] - 50, peaks_array[i][1] + 50, 1)
        #    temp_yrange = fgx[temp_xrange]
        #
        #elif next_peak_dif < 50:
        #    if peaks_array[i][1] < 50:
        #        temp_xrange = np.arange(0, peaks_array[i][1] + 25, 1)
        #    elif peaks_array[i][1] > len(pixels) - 50:
        #        temp_xrange = np.arange(peaks_array[i][1] - 50, len(pixels))
        #    elif peaks_array[i][1] >= 50:
        #        temp_xrange = np.arange(peaks_array[i][1] - 50, peaks_array[i][1] + 25, 1)
        #    temp_yrange = fgx[temp_xrange]
        #
        #elif previous_peak_dif < 50: 
        #    if peaks_array[i][1] < 50:
        #        temp_xrange = np.arange(0, peaks_array[i][1] + 50, 1)
        #    elif peaks_array[i][1] > len(pixels) - 50:
        #        temp_xrange = np.arange(peaks_array[i][1] - 25, len(pixels))
        #    elif peaks_array[i][1] >= 50:
        #        temp_xrange = np.arange(peaks_array[i][1] - 25, peaks_array[i][1] + 50, 1)
        #    temp_yrange = fgx[temp_xrange]  
        #
        #popt, pcov = curve_fit(gauss, temp_xrange, temp_yrange, [peaks_array[i][0], peaks_array[i][1], peaks_array[i][2], peaks_array[i][3]])
        #plt.plot(temp_xrange, gauss(temp_xrange, popt[0], popt[1], popt[2], popt[3]), 'r')
        
test_array_Cd650 = [[8E6, 150, 30, 3E5], [1.6E7, 270, 30, 3E5], [1.6E7, 520, 30, 3E5]]
test_array_Cd150 = [[8E5, 20, 40, 6E5], [8E5, 75, 20, 6E5], [2E6, 150, 30, 6E5], [1.6E6, 190, 30, 6E5], [9E5, 200, 30, 6E5], [1E6, 310, 20, 6E5], [1E6, 320, 30, 6E5], [1.6E6, 470, 20, 6E5], [1E6, 560, 30, 6E5], [8E5, 750, 30, 6E5]]
gauss_fits(test_array_Cd150)

plt.show()


# sigma = sqrt(mu * t + sigma_0**2)
# for normal dist, FWHM = 2.35 * stddev