# 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_6.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, 'b')
#
#plt.subplot(2, 1, 2)
yerr = np.sqrt(fgx)
plt.errorbar(pixels, fgx, yerr, 0, 'b.', ecolor='g')

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), 1)
        #else:
        #    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')
        
        next_peak_dif = 0
        previous_peak_dif = 0
        if i == 0:
            next_peak_dif = peaks_array[i + 1][1] - peaks_array[i][1]
            previous_peak_dif = 100
        elif i == len(peaks_array) -1:
            next_peak_dif = 100
            previous_peak_dif = peaks_array[i][1] - peaks_array[i - 1][1]
        else:
            next_peak_dif = peaks_array[i + 1][1] - peaks_array[i][1]
            previous_peak_dif = peaks_array[i][1] - peaks_array[i - 1][1]
        
        temp_xrange = []
        temp_yrange = []
        if next_peak_dif < 50 & previous_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] - 25, len(pixels), 1)
            else:
                temp_xrange = np.arange(peaks_array[i][1] - 25, peaks_array[i][1] + 25, 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), 1)
            else:
                temp_xrange = np.arange(peaks_array[i][1] - 50, peaks_array[i][1] + 30, 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), 1)
            else:
                temp_xrange = np.arange(peaks_array[i][1] - 25, peaks_array[i][1] + 50, 1)
            temp_yrange = fgx[temp_xrange]  
            
        else:
            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), 1)
            else:
                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]])
        
        chi_squared = 0.
        for j in range(0, len(temp_yrange)):
            chi_squared += (temp_yrange[j] - gauss(temp_xrange, popt[0], popt[1], popt[2], popt[3])[j])**2 / temp_yrange[j]
        #print chi_squared
        
        plt.plot(temp_xrange, gauss(temp_xrange, popt[0], popt[1], popt[2], popt[3]), 'r')
        plt.text(peaks_array[i][1] - 70, temp_yrange.max() + 0.6E6, chi_squared)
        plt.title("Cd Spectrum at 06.50mm with Gaussian fits and $\chi ^2$ values \n")
        plt.xlabel("Pixel Values")
        plt.ylabel(r"$N_{\rm{PE}}$") 
        
        print pcov
        
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_Cd650)

plt.show()


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

# this is 'instruental response function'
# measure shape
# weighted average -> xbar

# error bars -> first guess sqrt
# zoomed in to inspect
# chi squared
# include multpile params ?
# pcov is covariance matrix -> sqrt diag elements ?
# Anisha's code could lead to good estimate for cut off level