# -*- coding: iso-8859-1 -*-

from numpy import *
from scipy.signal import *
from scipy.optimize import *
from os import path, rename

def result():

    the_experiment = None	# current experiment's name
    var_key = ''
    
    measurements = {'satrec2_experiment': MeasurementResult('Saturation Recovery'),
                  'solidecho_experiment': MeasurementResult('Solid Echo'),
                     'spinal_experiment': MeasurementResult('Spin Alignment'),
                     'zeeman_experiment': MeasurementResult('Zeeman Order')}
                               
    measurement_ranges = {'satrec2_experiment': [0.5e-6, 4.5e-6],
                        'solidecho_experiment': [0.5e-6, 4.5e-6],
                           'spinal_experiment': [0.5e-6, 4.5e-6],
                           'zeeman_experiment': [0.5e-6, 4.5e-6]}
    measurement_ranging = True
                             
    # loop over the incoming results: 
    for timesignal in results:
        if not isinstance(timesignal,ADC_Result): 
            continue     

        # read experiment parameters:
        pars = timesignal.get_description_dictionary()
                    
        # keep track of the actual experiment's name:
        if the_experiment != pars.get('PROG'):
           the_experiment = pars.get('PROG')
           suffix = ''	# output file name's suffix
           counter = 1

        # ---------------- digital filter ------------------ 
             
        # get actual sampling rate of timesignal:
        sampling_rate = timesignal.get_sampling_rate()

        # get user-defined spectrum width:
        spec_width = pars['SW']
        
        # specify cutoff frequency, in relative units:
        cutoff = spec_width / sampling_rate
                    
        # number of filter's coefficients:
        numtaps = 29
                
        if cutoff < 1: # otherwise no filter applied
            
            # use firwin to create a lowpass FIR filter:
            fir_coeff = firwin(numtaps, cutoff)
            
            # downsize x according to user-defined spectral window:
            skip = int(sampling_rate / spec_width)
            timesignal.x = timesignal.x[::skip]
        
            for i in range(2):
                # apply the filter to ith channel:
                timesignal.y[i] = lfilter(fir_coeff, 1.0, timesignal.y[i])

                # zeroize first N-1 "corrupted" samples:
                timesignal.y[i][:numtaps-1] = 0.0

                # circular left shift of y:
                timesignal.y[i] = roll(timesignal.y[i], -(numtaps-1))    

                # downsize y to user-defined number of samples (SI):
                timesignal.y[i] = timesignal.y[i][::skip]

            # update the sampling_rate attribute of the signal's:
            timesignal.set_sampling_rate(spec_width)

        # ----------------------------------------------------

        # rotate timesignal according to current receiver's phase:
        timesignal.phase(pars['rec_phase'])

        # provide timesignal to the display tab:
        data['Current scan'] = timesignal

        # accumulate...
        if not locals().get('accu'):
            accu = Accumulation()

        # skip dummy scans, if any:
        if pars['run'] < 0: continue

        # add up:
        accu += timesignal

        # provide accumulation to the display tab:
        data['Accumulation'] = accu

        # check how many scans are done:
        if accu.n == pars['NS']: # accumulation is complete
           
	    # make a copy:
            echo = accu + 0

            # compute the signal's phase:
            #phi0 = arctan2(accu.y[1][0], accu.y[0][0]) * 180 / pi
            #if not locals().get('ref'): ref = phi0   
            #print 'phi0 = ', phi0

            # rotate the signal to maximize Re (optional):
            #echo.phase(-phi0)

            # check whether it is an arrayed experiment: 
            var_key = pars.get('VAR_PAR')
            if var_key:
                # get variable parameter's value:
                var_value = pars.get(var_key)	

                # provide signal recorded with this var_value to the display tab:
                data['Accumulation'+"/"+var_key+"=%e"%(var_value)] = accu

		# estimate noise level:
                if not locals().get('noise'):
                    n = int(0.1*accu.x.size)
                    noise = 3*std(accu.y[0][-n-29:-30])  				
					
                # measure signal intensity vs. var_value:	
                if the_experiment in measurements.keys():
				
                    # option a: the sum of samples within the given range:
                    if measurement_ranging == True:
                        [start, stop] = echo.get_sampling_rate() * array(measurement_ranges[the_experiment])
                        measurements[the_experiment][var_value] = sum(echo.y[0][int(start):int(stop)]) 
						
                    # option b: the sum of all samples above noise: 
                    else:
                        measurements[the_experiment][var_value] = sum(echo.y[0][echo.y[0]>noise])
                    
                    # store a measurement:
                    data[measurements[the_experiment].get_title()] = measurements[the_experiment]
                    
                    # update the file name suffix:
                    suffix = '_' + str(counter)
                    counter += 1
					
                else: 
                    print "Cannot recognize experiment: accumulate without measuring"				

            # save accu if required:   
            outfile = pars.get('OUTFILE')
            if outfile: 
                datadir = pars.get('DATADIR')

                # write raw data in Simpson format:
                filename = datadir+outfile+suffix+'.dat'
                if path.exists(filename):
                    rename(filename, datadir+'~'+outfile+suffix+'.dat')
                accu.write_to_simpson(filename)

               # write parameters in a text file:
                filename = datadir+outfile+suffix+'.par' 
                if path.exists(filename):
                    rename(filename, datadir+'~'+outfile+suffix+'.par')

                fileobject = open(filename, 'w')
                for key in sorted(pars.iterkeys()):
                    if key=='run': continue
                    if key=='rec_phase': continue
                    fileobject.write('%s%s%s'%(key,'=', pars[key]))
                    fileobject.write('\n')
                fileobject.close()

            # reset accumulation:
            del accu

    if var_key == 'TAU' or var_key == 'D2':
       # mono-exponential saturation-recovery fit:
        try:
            xdata = measurements['satrec2_experiment'].get_xdata()
            ydata = measurements['satrec2_experiment'].get_ydata()
            [amplitude, rate, offset] = fitfunc_recovery(xdata, ydata)
            print 'Mono-exponential fit to saturation-recovery data:'
            print '%s%02g' % ('Amplitude = ', amplitude)
            print '%s%02g' % ('T1 [s] = ', 1./rate)

            # update display for fit:      
            measurements['satrec2_experiment'].y = func_recovery([amplitude, rate, offset], xdata)
            data[measurements['satrec2_experiment'].get_title()] = measurements['satrec2_experiment']
        except:
            pass
		
        # KWW fit to spin-alignment echoes:
        try:
            xdata = measurements['spinal_experiment'].get_xdata()
            ydata = measurements['spinal_experiment'].get_ydata()
            [amplitude, rate, beta] = fitfunc_kww(xdata, ydata)
            print 'KWW fit to spin-alignment echoes:'
            print '%s%02g' % ('Amplitude = ', amplitude)
            print '%s%02g' % ('T2 [s] = ', 1./rate)
            print '%s%01g' % ('Beta = ', beta)

            # update display for the fit:
            measurements['spinal_experiment'].y = func_kww([amplitude, rate, beta], xdata)
            data[measurements['spinal_experiment'].get_title()] = measurements['spinal_experiment']
        except:
            pass
		
        # KWW fit to Zeeman-order echoes:
        try:
            xdata = measurements['zeeman_experiment'].get_xdata()
            ydata = measurements['zeeman_experiment'].get_ydata()
            [amplitude, rate, beta] = fitfunc_kww(xdata, ydata)
            print 'KWW fit to Zeeman-order echoes:'
            print '%s%02g' % ('Amplitude = ', amplitude)
            print '%s%02g' % ('T2 [s] = ', 1./rate)
            print '%s%01g' % ('Beta = ', beta)

            # update display for the fit:
            measurements['zeeman_experiment'].y = func_kww([amplitude, rate, beta], xdata)
            data[measurements['zeeman_experiment'].get_title()] = measurements['zeeman_experiment']
        except:
            pass	

# the fitting procedure for satrec_experiment:
def fitfunc_recovery(xdata, ydata):

    # initialize variable parameters:
    try: 
        # solve Az = b:
        A = array((ones(xdata.size/2), xdata[0:xdata.size/2]))    
        b = log(abs(mean(ydata[-2:]) - ydata[0:xdata.size/2]))     
        z = linalg.lstsq(transpose(A), b)
        amplitude = exp(z[0][0])
        rate = -z[0][1]
    except:
        amplitude = abs(ydata[-1] - ydata[0])
        rate = 1./(xdata[-1] - xdata[0])
    offset = min(ydata)
    p0 = [amplitude, rate, offset]

    # run least-squares optimization:
    plsq = leastsq(residuals_recovery, p0, args=(xdata, ydata))

    return plsq[0] # best-fit parameters

def residuals_recovery(p, xdata, ydata):
    return ydata - func_recovery(p, xdata)

# here is the function to fit
def func_recovery(p, xdata):
    return p[0]*(1-exp(-p[1]*xdata)) + p[2]

			
# the fitting procedure for spinal_experiment and zeeman_experiment:
def fitfunc_kww(xdata, ydata): 

    # initialize variable parameters:
    try: 
        # solve Az = b:
        A = array((ones(xdata.size/2), xdata[0:xdata.size/2]))    
        b = log(abs(ydata[0:xdata.size/2]))     
        z = linalg.lstsq(transpose(A), b)
        amplitude = exp(z[0][0])
        rate = -z[0][1]
    except:
        amplitude = abs(ydata[0])
        rate = 1./(xdata[-1] - xdata[0])
    beta = 1.
    p0 = [amplitude, rate, beta]

    # run least-squares optimization:
    plsq = leastsq(residuals_kww, p0, args=(xdata, ydata))

    return plsq[0] # best-fit parameters

def residuals_kww(p, xdata, ydata):
    return ydata - func_kww(p, xdata)

# here is the function to fit:
def func_kww(p, xdata):
    return p[0]*exp(-(p[1]*xdata)**p[2])


pass