import math

from mathutils import fit_polynomial, PolynomialFunction
from plotutils import plot_line, plot_function
from ch.psi.pshell.swing.Shell import STDOUT_COLOR
import org.apache.commons.math3.stat.correlation.PearsonsCorrelation as PearsonsCorrelation

if get_context().source == CommandSource.ui:
    #dx = "SINEG01-RLLE-REF10:SIG-PHASE-AVG"
    #dy = "SINEG01-RLLE-REF20:SIG-PHASE-AVG" 
    dx = "SINEG01-RGUN-PUP10:SIG-AMPLT-AVG"
    dy = "SINEG01-RGUN-PUP20:SIG-AMPLT-AV"
    interval = 0.10
    window = 40
    p = plot(None)[0]    
    bs = True

for s in p.getAllSeries():
    p.removeSeries(s)    

_stream = None

instances = []
def _get_device(d):
    global _stream
    if isinstance(d, basestring):    
        try:
            if get_device(d) is not None: 
                return get_device(d)      
            r = eval(d)
            print d
            if isinstance(r, Device): 
                return r
        except:
            pass
        if bs == True:
            if _stream == None:
                _stream = Stream("corr_stream", dispatcher)
                instances.append(_stream)
            _stream.addScalar(d, d, int(interval*100), 0)
        else:
            d =  Channel(d)    
            instances.append(d)
    return d
dx = _get_device(dx)
dy = _get_device(dy)
        
try:
    if _stream != None:
        _stream.initialize()
        _stream.start()
    p.addSeries(LinePlotSeries("Data")) 
    sd=p.getSeries(0)
    sd.setLinesVisible(False)
    
    sd.setPointSize(4)
    
    if globals().has_key("marker"):
        p.removeMarker(marker)
    
    marker=None
    corr = None
    
    
    while(True):
        #Sample and plot data
        if bs == True:
            _stream.waitValueNot(_stream.take(), 10000)
            (x,y) = _stream.take().values
        else:
            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:
            x1, x2 = min(ax), max(ax)  
            res = (x2-x1)/100     
            if x1!=x2:
                #Display correlation
                corr= PearsonsCorrelation().correlation(to_array(ax,'d'), to_array(ay,'d'))        
                s = "Correlation=" + str(round(corr,4))
                #print s
                if marker is not None: 
                    p.removeMarker(marker)
                marker = p.addMarker(x2+res, p.AxisId.X, s,  p.getBackground())
                marker.setLabelPaint(STDOUT_COLOR)
                
                #Calculate, print and plot linear fit
                pars_lin = (a0,a1) = fit_polynomial(ay, ax, 1)
                #print "Fit lin  a1:" , a1, " a0:",a0        
                y1 = poly(x1, pars_lin)
                y2 = poly(x2, pars_lin)
                plot_line(p, x1, y1, x2, y2, width = 2, color = Color.BLUE, 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)
                ax = frange(x1, x2, res, True)
                plot_function(p, fitted_quad_function, "Fit Quadratic", ax, color=Color.GREEN)
        if bs != True:
            time.sleep(interval)
finally:
    for dev in instances:
        dev.close()