import numpy as _np
import pyfits as _pyfits
import scipy.optimize as _spo
import matplotlib.pyplot as _plt
from scipy.ndimage import median_filter as _median_filter

'''
def test(filename="./He_Rotation/He_Rotation_0350.fit") :
    d = prepare_data(filename);
    r = fit_ian(d);
    gauss_fit_plot(d, r)
'''

def get_raw_data(filename) : 
    f = _pyfits.open(filename)
    h = f[0].header
    d = f[0].data * h.get('EGAIN')
    return d

def raw_image_display(filename) : 
    d = get_raw_data(filename)
    _plt.clf()
    _plt.imshow(d)
    _plt.title("Helium line at 4.20mm")
    _plt.xlabel("Pixels")
    _plt.ylabel("Pixels")
    
def raw_spectrum_display(filename):
    d = get_raw_data(filename)
    fgx = d.sum(0)
    fgy = d.sum(1)
    pixels = _np.arange(0, len(fgx), 1)
    _plt.plot(pixels, fgx)
    _plt.title("Profile of Helium peak at 7.80mm (388nm)")
    _plt.xlabel("Pixel position")
    _plt.ylabel(r"$N_{\rm{PE}}$")
    
'''
def hotpixel_removal(imagedata) :
    return 
''' 

def project(filename):
    d = get_raw_data(filename)
    px = d.sum(0)
    py = d.sum(1)
    return [px, py]

'''
def prepare_data(filename) :
    d = get_raw_data(filename)
    # remove hot pixels 
    px = d.sum(0)
    return px
'''

'''
def locate_peaks(data) : 
    pass
'''

'''    
def stitch(filelist) :
'''

'''
def pixel_to_wavelength(filelist):
'''

def weighted_av(x, y):
    return sum(x) / sum(y)
    
def standev(x, xbar, n):
    return _np.sqrt(_np.sum((x - xbar)**2) / n)

def multi_gauss(x,*p): # this is a gauss fit for multiple peaks rather than a single peak.
    y=_np.zeros_like(x)
    for i in range(0,len(p),4):
        a=p[i]
        x0=p[i+1]
        sigma=p[i+2]
        bg=p[i+3]
        y= y+ a / (_np.sqrt(2.0*_np.pi)*sigma) * _np.exp(-(x - x0)**2 / (2 * sigma**2))+bg
    return y
    
def get_peak_var(filename, cutoff, length): # cutoff is the cut off value to consider the values a peak. length is the nth adjacent value taken away from the middle value when calculating the second differential
    p  = project(filename)
    px = p[0]
    x  = _np.arange(0, len(px) - (length + 1), 1) # list of values to be used as the pixel values
    bgsubmaxnorm = (px[0:-length] - px[0:-length].mean()) / px[0:-length].max() #finds the normlised photoelectrons with the subtracted background
    xdiff = _np.diff(px) # obtain first differential
    sec = []
    for i in range(0, len(xdiff) - length):
        peak_val = xdiff[i + length] - xdiff[i - length]
        sec.append(peak_val)        
    secdiff = _np.asarray(sec)  # gets the second differential but instead of subtracting adjacent values it subtracts every nth(lenght) value and converts it to array      

    secdiffynorm          = secdiff / secdiff.max()
    #filtered_secdiffynorm = _median_filter(secdiffynorm, size = 4) # helps finding better estimation for peak
    x_peak_range       = x[secdiffynorm < cutoff] # > for absroption, < for emission
    return [x, secdiffynorm, x_peak_range, bgsubmaxnorm]

def peak_range_peakfinder(filename, cutoff, length):
    peak_var      = get_peak_var(filename, cutoff, length)
    x_peak_range  = peak_var[2]
    bgsubmaxnorm  = peak_var[3]
    list_of_lists = [] 
    sub_list      = []
    sub_list.append(x_peak_range[0])  #takes the first value of xmer and adds it to the sublist
    
    for i in range(0, len(x_peak_range) - 1): 
        if x_peak_range[i] +  1 == x_peak_range[i + 1]:  #if the ith value is 1 less than ith+1 value then it adds it to the sub_list 
            sub_list.append(x_peak_range[i + 1]) 
        else:  #if that is not the case then it adds the sub_list to another list then makes a new empty sub_list and starts the process again.
            if len(sub_list) > 0:
                list_of_lists.append(sub_list)
                sub_list = [] 
                sub_list.append(x_peak_range[i + 1])
    
    list_of_lists.append(sub_list) # all of the sub_list are added to the main list so it has groups of x values which have higher/lower intensity than the cutoff value
    
    
    xy = []
    y  = []

    for i in list_of_lists:
        multiple = bgsubmaxnorm[i] * i
        xy.append(multiple)
        y.append(bgsubmaxnorm[i])
    return [list_of_lists, xy, y]

def Peaklocator(filename, cutoff, length, value):
    list_of_lists = peak_range_peakfinder(filename, cutoff, length)[0]
    xy = peak_range_peakfinder(filename, cutoff, length)[1]
    y = peak_range_peakfinder(filename,cutoff,length)[2]
    px = project(filename)[0]
    
    '''
    for i in range(0,len(xy)):
        if len(xy[i]) > value: 
            print 'xlocation->', int(weighted_av(xy[i], y[i]))
            print 'intensity->', px[int(weighted_avxy[i], y[i]))]
    '''
    peaklocation=[]
    intensity=[]
    for i in range(0,len(xy)):
        if len(xy[i])>value: 
            ploc= int(weighted_av(xy[i],y[i])) # uses the weighted average equation to find where the location of the peak and its intensity at that point.
            #print 'xlocation->', ploc
            #print 'intensity->', px[int(weigave(xy[i],y[i]))] # intensity at the location of estimated peak
            peaklocation.append(ploc)
            intensity.append(px[ploc])
                       
    num = [] # numerator of the weighted average
    denom = [] # denominator of the weighted average
    for i in list_of_lists:
        multiple = px[i] * (_np.square(i - _np.mean(i)))
        num.append(multiple)
        denom.append(px[i])

    width = []
    for i in range(0, len(num)):
        if len(xy[i]) > value:
            x = weighted_av(num[i], denom[i])
            width.append(x)   
   
    
    return [width,peaklocation,intensity]

def plot_vertline_peaks(filename, cutoff, length, value):
    xy=peak_range_peakfinder(filename, cutoff, length)[1]
    y=peak_range_peakfinder(filename,cutoff,length)[2]
    px=project(filename)[0]
    
    _plt.title('Plot of spectrum of Vega at 06.50mm with located peaks') # change title according to plot
    _plt.xlabel('Pixel values')
    _plt.ylabel('Photoelectrons')
    _plt.plot(px, linewidth='2.0', label='vega spectral line')
    a='Estimated peak location ' #this is for multiple peaks so it says 'peak 1', 'peak 2' and so on

    #a=0
    for i in range(0,len(xy)):
        if len(xy[i])>2: # this statement is here to get rid of any peaks that could be due to hot pixels as their "width" is too small.
            #a+=1
            b=str(i+1)
            c=a+b
            _plt.axvline(weighted_av(xy[i],y[i]), linewidth='2.0', color='red',label=c)
           
    _plt.legend()
    _plt.show()


#here the filenames corresponds to the calibration files -> Calib_He_0600.fit
def multi_gaussfit(filename,cutoff,length,value):
    ydata=project(filename)[0] #get ydata
    xdata=_np.arange(0,len(ydata),1) # create the xdata for xaxis
    Peak=_np.asarray(Peaklocator(filename,cutoff,length,value)[1]) # get unfitted peaklocation to use as guess paramters
    intensity=_np.asarray(Peaklocator(filename,cutoff,length,value)[2]) # get unfitted intensity at the peaklocation to use as guess parameters
    sigma=[10,6,8,9] # guess sigma
    bg=[10000,10000,10000,10000]#guess the background level

    guess=[] #empty aray
    for i in range(0,len(Peak)):
        guess.append(intensity[i]) #makes the first value in the array the intensity value and so on
        guess.append(Peak[i])
        guess.append(sigma[i])
        guess.append(bg[i])
        
    popt,pcov=_spo.curve_fit(multi_gauss, xdata, ydata, p0=guess) # gives back fitted values for peaklocation, intesnity, sigma and background
    fit=multi_gauss(xdata,*popt)
    #_plt.figure()
    #_plt.plot(xdata,ydata)#plots all the fitted curves in one canvas
    #_plt.plot(xdata,fit,'r-')
    #_plt.show()
    
    newpopt= [popt[x:x+4] for x in range(0, len(popt),4)] # seperates the returned values from the curvefit to 4 length arrays: intensity, peaklocation, sigma, bg
    intens=[]
    Plocs=[]
    sig=[]
    bgs=[]
   
    for i in range(0,len(newpopt)):#seperate the parameters and form their own arrays to be used
        intens.append(newpopt[i][0])
        Plocs.append(newpopt[i][1])
        sig.append(newpopt[i][2])
        bgs.append(newpopt[i][3])
        
    return [intens,Plocs,sig,bgs]
  
#filename corresponds loop_mult_files is helium.txt      
def loop_mult_files(filename,cutoff,length,value): # just a function to loop over the multi_gauss fit function to obtain intensity, peaklocation, bg, sigma, wavelength for different micrometer position.
    f=open(filename)
    Peak_locs=[]
    micro_pos=[]
    wavelength=[]
    intensity=[]
    for l in f: # loop through the files and extract values
        l=l.strip()
        sl=l.split()
        filelist=sl[0]
        peak=multi_gaussfit(filelist,cutoff,length,value)[1] 
        Peak_locs.append(peak)
        micro_pos.append(float(sl[1]))
        wavesl=sl[2].split(",")
        wave=map(float,wavesl)
        intsl=multi_gaussfit(filelist,cutoff,length,value)[0]
        intensity.append(intsl)
        wavelength.append(wave)
            
    return[Peak_locs, micro_pos,wavelength,intensity] #get better peak locations and intensity.

def lambdafunc((m,p),I,a,b,c,d,e,p0,m0):

    a=I+a*(p-p0)-b*(m-m0)+c*(p-p0)**2+d*(m-m0)**2-e*(p-p0)*(m-m0)
    return a.ravel()
   
def pixel_to_wavelength(filename,filelist,cutoff,length,value,mpos):
    micro_pos=loop_mult_files(filename,cutoff,length,value)[1]
    Xposition=loop_mult_files(filename,cutoff,length,value)[0]
    wave=loop_mult_files(filename,cutoff,length,value)[2]
    Peakloc=multi_gaussfit(filelist,cutoff,length,value)[1]
    px=project(filelist)[0]
    xdata=_np.arange(0,len(px),1)

    print wave
    micro=[]
    for i in range(0,len(Xposition)):
        for j in range(0,len(Xposition[i])):
            micro.append(micro_pos[i])
    
    Xpos=[item for sublist in Xposition for item in sublist]
    wavelength=[item for sublist in wave for item in sublist]
    
    
    Part_wave=[]
    
    for i in range(0,len(wavelength)):
        if micro[i]==mpos:
            Part_wave.append(wavelength[i])
            
    print Part_wave
    print Peakloc
    peaksq=_np.square(Peakloc)
    
    grad,inter=_np.polyfit(Peakloc,Part_wave,1)
    fit=_np.polyfit(Peakloc,Part_wave,1)
    grad2=_np.polyfit(peaksq,Part_wave,2)
    fit_fn=_np.poly1d(fit)
    #print fit_fn
    best_fit=fit_fn(xdata)
    _plt.figure(1)
    _plt.plot(Peakloc,Part_wave,'o',label='Relation between wavelength and pixel at micrometer of %s mm ' %mpos, linewidth=5)
    _plt.plot(xdata,best_fit,'r-', label='Line of best fit',linewidth=2)
    _plt.title('Relation between pixel values and ' r'wavelength,$\lambda$ (nm)', size=22)
    _plt.xlabel('Pixel Position',size=22)
    _plt.ylabel(r'Wavelength,$\lambda$ (nm)',size=22)
    _plt.legend(loc='best')
    
    
    print inter
    print grad
    print grad2
    guess=[inter,grad,0.2,grad2[1],2e-3,0.8,365,4.2]
    popt,pcov=_spo.curve_fit(lambdafunc,(micro,Xpos),wavelength,p0=guess)
    print popt
    data_fitted=lambdafunc((mpos,xdata),*popt)
    #vertline=plot_vertline_peaks(filelist, cutoff, length, value)
    _plt.figure(2)
    _plt.plot(data_fitted,px,linewidth='3',label='Projection of helium spectral lines at %s' %mpos)
    _plt.xlim(data_fitted[0],data_fitted[764])
    _plt.title('Plot of wavelength against number of photoelectrons for Helium spectrum', size=22)
    _plt.xlabel(r'Wavelength,$\lambda$ (nm)',size=18)
    _plt.ylabel('Number of photoelectrons', size=18)
    _plt.legend()
    _plt.show()
    
'''
def pixel_to_wavelength(w1,w2,w3,w4,filename,cutoff,length,value):# for Calib_He_0600 the wavelengths are 447,471,492,501 and cutoff length and value are -0.05,1,2
    peaklocation=Peaklocator(filename, cutoff, length, value)[1]
    #peaksq=_np.square(peaklocation)
    #px=prepare_data(filename)
    #xint=_np.arange(0,765,1)
    True_wavelength=[w1,w2,w3,w4]
    #params=_spo.curve_fit(lambdafunc,peaklocation,True_wavelength, p0=[0.1,0.1,-1e-8,-2e-8,0.8,5,365,6,4.2])
    popt,pcov=_spo.curve_fit(lambdafunc,peaklocation,True_wavelength, p0=[446,0.1,0.1,-1e-8,-2e-8,0.8,365,6,4.2])

    print popt,pcov


def lambdafunc((m,p),I,a,b,c,d,e,p0,m0):
    return I+a*(p-p0)-b*(m-m0)+c*(p-p0)**2+d*(m-m0)**2-e*(p-p0)*(m-m0)


def pixel_to_wavelength(filename,filelist,cutoff,length,value):
    micro_pos=loop_mult_files(filename,cutoff,length,value)[1]
    Xposition=loop_mult_files(filename,cutoff,length,value)[0]
    wave=loop_mult_files(filename,cutoff,length,value)[2]
    px=project(filelist)[0]
    xdata=_np.arange(0,len(px),1) 
    micro=[]
    for i in range(0,len(Xposition)): # this is to make it length of micro the same as the xpostion and wavelength
        for j in range(0,len(Xposition[i])):
            micro.append(micro_pos[i])
    
    Xpos=[item for sublist in Xposition for item in sublist] # flatten out the Xpos so instead of [[123,143],[156,160]] you get [123,143,156,160] 
    wavelength=[item for sublist in wave for item in sublist]
    guess=[446,0.1,0.2,-1e-8,-2e-8,0.8,365,4.2] # random guess for the parameters
    popt,pcov=_spo.curve_fit(lambdafunc,(micro,Xpos),wavelength,p0=guess)
    print popt
    data_fitted=lambdafunc((xdata,px),*popt)#this is random attempt to plot the lambdafunction
    _plt.plot(data_fitted)
    _plt.show()

'''
    
'''
def chi_squared
'''

def single_gauss_plot(filename, amp, mu, sigma, bg):
    d = get_raw_data(filename)
    p = project(filename)
    x = _np.array(d[0])
    y = _np.array(d[1])
    px = p[0]
    py = p[1]
    popt, pcov = _spo.curve_fit(gauss, x, y, [amp, mu, sigma, bg])
    poptcov = [popt, pcov]
    yfit = gauss(x, poptcov[0][0], poptcov[0][1], poptcov[0][2], poptcov[0][3])
    print poptcov
    _plt.clf()
    _plt.grid(True)
    _plt.plot(x, px, 'b.')
    _plt.plot(x, yfit, 'r')
    
def gauss(x, a, x0, sigma, bg):
    return a/(_np.sqrt(2.0*_np.pi)*sigma) * _np.exp(-(x - x0)**2 / (2 * sigma**2)) + bg

def fit_gaussian(y, amp=1000000, mu=500, sig=30, bg=100000) : 
    x = _np.arange(0, len(y), 1)
    popt, pcov = _spo.curve_fit(gauss, x, y, [amp, mu, sig, bg])
    return [popt, pcov]
    
def gauss_fit_plot(d, poptcov) : 
    x    = _np.arange(0,len(d),1)
    y    = d 
    yfit = gauss(x, poptcov[0][0], poptcov[0][1], poptcov[0][2], poptcov[0][3])
    _plt.clf()
    _plt.grid(True)
    #_plt.plot(d, 'b.')
    #_plt.plot(yfit)
    _plt.show()

def gauss_fit_loop(filelist) :
    f = open(filelist)
    val_array = []
    sig_array = []
    for l in f : 
        l   = l.strip()
        sl  = l.split() 
        val = float(sl[0])
        filename = sl[1]
        d = project(filename)
        r = fit_gaussian(d[0])
        val_array.append(val)
        sig_array.append(abs(r[0][2]))
    #_plt.grid(True)
    #_plt.plot(val_array, sig_array, 'b.')
    return [val_array, sig_array]

def hyperbole(x, a, b, mu):
    return -b * _np.sqrt( ((x-mu)/a)**2 + 1.0)
    
def fit_hyperbole(x, y, a = 1.0, b = 4.0, mu = 21.0):
    popt, pcov = _spo.curve_fit(hyperbole, x, y, [a, b, mu])
    return [popt, pcov]

def hyperbole_fit_plot(d, poptcov) :
    x    = d[0]
    y    = d[1]
    yfit = hyperbole(x, poptcov[0][0], poptcov[0][1], poptcov[0][2])
    #Extract x value
    _plt.clf()
    _plt.grid(True)
    _plt.plot(x, y, 'b.', label='Sigma values')
    _plt.plot(x, yfit, 'g', label='Hyperbolic fit')
    #_plt.title("Sigma of the Gaussian fit to the peak \n as a function of the lens-camera distance")
    #_plt.xlabel("Micrometer distance (mm)")
    #_plt.ylabel(r"$\sigma$")
    _plt.legend(loc='upper left')
    
def hyperbole_fit_loop(filelist):
    d       = _np.array(gauss_fit_loop(filelist))
    x       = _np.array(d[0])
    y       = _np.array(d[1])
    poptcov = fit_hyperbole(x, y, 1, 4.0, 21.0)
    yfit    = hyperbole_fit_plot(d, poptcov)
    
def cosineF(x, d, sm, k, phi):
    return d * _np.cos(k * x + phi) + sm

def fit_cosineF(x, y, d = 1.0, sm = 4.0, k = 0.25, phi = 1.0):
    popt, pcov = _spo.curve_fit(cosineF, x, y, [d, sm, k, phi])
    return [popt, pcov]
    
def cosineF_fit_plot(d, poptcov):
    x    = d[0]
    y    = d[1]
    yfit = cosineF(x, poptcov[0][0], poptcov[0][1], poptcov[0][2], poptcov[0][3])
    #Extract x value
    _plt.clf()
    _plt.grid(True)
    _plt.plot(x, y, 'b.', label='Sigma values')
    _plt.plot(x, yfit, 'g', label='Cosine function fit')
    #_plt.title("Sigma of the Gaussian fit to the peak \n as a function of the camera rotation")
    #_plt.xlabel("Micrometer (mm)")
    #_plt.ylabel(r"$\sigma$")
    _plt.legend(loc='upper left')
    
def cosineF_fit_loop(filelist):
    d       = _np.array(gauss_fit_loop(filelist))
    x       = _np.array(d[0])
    y       = _np.array(d[1])
    poptcov = fit_cosineF(x, y, 1.0, 4.0, 0.25, 1.0)
    yfit    = cosineF_fit_plot(d, poptcov)
    
'''
def profile_loop(filelist):
    f = open(filelist)
    val_array = []
    for l in f:
        l = l.strip()
        sl = l.split() 
        val = float(sl[0])
        val_array.append(val)
        filename  = sl[1]
        d = project(filename)
        x = _np.array(d[0])
        y = _np.array(d[1])
        ymax_index = y.argmax()
        xmax_index = ymax_index
        temp_xrange = _np.arange(xmax_index - 50, xmax_index + 50, 1)
        temp_yrange = y[temp_xrange]
        _plt.clf()
        _plt.grid(True)
        _plt.plot(temp_xrange, temp_yrange)
'''

def asymmetry_profile_loop(filelist):
    d = _np.array(gauss_fit_loop(filelist))
    x = _np.array(d[0])
    y = _np.array(d[1])
    _plt.plot(x, y, 'b.')
    

'''
To do: 

Code
Lamps
Lambda func - alpha
Sigma func
Observation
Stitch (?)
Hot pixels
'''