from mathutils import estimate_peak_indexes, fit_gaussians, create_fit_point_list, Gaussian
import java.awt.Color as Color
 
import mathutils
mathutils.MAX_ITERATIONS = 100000

def fit(ydata, xdata = None, draw_plot = True):
    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]
    print "Max index:" + str(index_max),
    print " x:" + str(max_x),
    print " y:" + str(max_y)
    
    if draw_plot:
        plots = plot([ydata],["data"],[xdata], title="Fit" )
        p = None if plots is None else plots[0]

    gaussians = fit_gaussians(ydata, xdata, [index_max,])
    if gaussians[0] is None:
        if draw_plot and (p is not None):
            p.addMarker(max_x, None, "Max="+str(round(max_x,4)), Color.GRAY)    
        print "Fitting error"
        return (None, None, None)
        
    (norm, mean, sigma) = gaussians[0]
    if draw_plot:
        fitted_gaussian_function = Gaussian(norm, mean, sigma)    
        scale_x = [float(min(xdata)), float(max(xdata)) ]
        points = max((len(xdata)+1), 100)
        resolution = (scale_x[1]-scale_x[0]) / points    
        fit_y = []
        fit_x = frange(scale_x[0],scale_x[1],resolution, True)
        for x in fit_x:
            fit_y.append(fitted_gaussian_function.value(x))            
        #Server
        if p is None:
            plot([ydata,fit_y],["data","fit"],[xdata,fit_x], title="Fit")             
            draw_plot = False
        else:
            p.addSeries(LinePlotSeries("fit"))
            p.getSeries(1).setData(fit_x, fit_y)            

    if abs(mean - xdata[index_max]) < abs((scale_x[0] + scale_x[1])/2):        
        if draw_plot:
            p.addMarker(mean, None, "Mean="+str(round(mean,4)), Color.MAGENTA.darker())
        print "Mean -> " +  str(mean)
        return (norm, mean, sigma)
    else:
        if draw_plot:
            p.addMarker(max_x, None, "Max="+str(round(max_x,4)), Color.GRAY)
        print "Invalid gaussian fit: " +  str(mean)
        return (None, None, None)