from mathutils import estimate_peak_indexes, fit_gaussians, create_fit_point_list, Gaussian
from mathutils import fit_polynomial,fit_gaussian, fit_harmonic, calculate_peaks
from mathutils import PolynomialFunction, Gaussian, HarmonicOscillator, LaguerreSolver
import math



def hfit(ydata, xdata = None):
    if xdata is None:
        xdata = frange(0, len(ydata), 1)

    max_y= max(ydata)    
    index_max = ydata.index(max_y)    
    max_x= xdata[index_max]
        
    start,end = min(xdata), max(xdata)
    (amplitude, angular_frequency, phase) = fit_harmonic(ydata, xdata)
    fitted_harmonic_function = HarmonicOscillator(amplitude, angular_frequency, phase)
    
    print "amplitude = ", amplitude
    print "angular frequency = ", angular_frequency
    print "phase = ", phase

    f = angular_frequency/ (2* math.pi)  
    print "frequency = ", f  

    resolution = 0.01
    fit_y = []
    for x in frange(start,end,resolution, True):
        fit_y.append(fitted_harmonic_function.value(x))
    fit_x = frange(start, end+resolution, resolution)

    p = plot(ydata,"data", xdata, title="HFit")[0]
    p.addSeries(LinePlotSeries("fit"))
    p.getSeries(1).setData(fit_x, fit_y)    

    #m = (phase + math.pi)/ angular_frequency
    m = -phase / angular_frequency
    if (m<start):
        m+=(1.0/f)
        
    if start <=m <=end:
        print "fit = ", m
        p.addMarker(m, None, "Fit="+str(round(m ,2)), Color.MAGENTA.darker())
        return (amplitude, angular_frequency, phase, True, m)
    else:
        print "max = ",max_x
        p.addMarker(max_x, None, "Max="+str(round(max_x ,2)), Color.MAGENTA.darker())        
        return (amplitude, angular_frequency, phase, False, max_x)
    

    

a= lscan(inp, (sin,out,arr), 0, 0.400, 10, 0.1)
hfit(a.getReadable(0),a.getPositions(0))