import math

from mathutils import fit_polynomial, PolynomialFunction
from plotutils import plot_line, plot_function

if get_context().source == CommandSource.ui:
    dx = "SINEG01-MSOL130:X"
    dy = "SINEG01-MSOL130:Y" 
    interval = 0.10
    window = 20
    p = plot(None, "Data")[0]

if isinstance(dx, basestring):    dx = Channel(dx)
if isinstance(dy, basestring):    dy = Channel(dy)

sd=p.getSeries(0)
#d.setLinesVisible(False)
sd.setPointSize(4)
p.addSeries(LinePlotSeries("Fit Quadratic"))
sfq=p.getSeries(1)

while(True):
    #Sample and plot data
    x=dx.read()  
    y=dy.read()
    sd.appendData(x, y)
    if len(sd.x) > window:
        #Remove First Element
        sd.token.remove(0)    
    ax = sd.x
    ay = sd.y
    if len(ax)>2:
        #Calculate, print and plot linear fit
        pars_lin = (a0,a1) = fit_polynomial(ay, ax, 1)
        print "Fit lin  a1:" , a1, " a0:",a0        
        x1 = min(ax); y1 = poly(x1, pars_lin)
        x2 = max(ax); y2 = poly(x2, pars_lin)
        plot_line(p, x1, y1, x2, y2, width = 2, color = Color.GREEN, name = "Fit Linear")

        #Calculate, print and plot quadratic fit
        pars_quad = (a0,a1,a2) = fit_polynomial(ay, ax, 2)
        print "Fit quad a2:" , a2, "a1:" , a1, " a0:",a0        
        fitted_quad_function = PolynomialFunction(pars_quad)
        ay=[]
        ax = frange(x1, x2, 1 if (x2==x1) else (x2-x1)/100, True)
        for x in ax:
            ay.append(fitted_quad_function.value(x))        
        plot_function(p, fitted_quad_function, "Fit Quadratic", ax, color=Color.BLUE)
        #sfq.setData(ax, ay)
            
    time.sleep(interval)