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

start_task("outupdate", 0.0, 0.0)

if get_exec_pars().source == CommandSource.ui:
    #dx = "SINEG01-RLLE-REF10:SIG-PHASE-AVG"
    #dy = "SINEG01-RLLE-REF20:SIG-PHASE-AVG"     
    #dx = "SINEG01-RGUN-PUP10:SIG-AMPLT-AVG 4"
    #dy = "SINEG01-RGUN-PUP20:SIG-AMPLT-AVG 4"

    #dx = "SINDI01-RKLY-DCP10:REF-AMPLT"
    #dy = "SINDI01-RKLY-DCP10:REF-PHASE"
            
    #dx = "SINDI01-RLLE-REF10:SIG-PHASE-AVG"
    #dy = "SINDI01-RLLE-REF20:SIG-PHASE-AVG"

    dx = "TESTIOC:TESTCALCOUT:Input"
    dx = "TESTIOC:TESTCALCOUT:Output"

    dy = "TESTIOC:TESTSINUS:SinCalc"
 
    #dx = "SINEG01-DICT215:B1_CHARGE"
    #dy = "SINEG01-DBPM314:Q1"     
    
    
    #dx=gsx.getReadback()
    #dy=gsy.getReadback()
    interval = 0.1
    window = 40
    p = plot(None)[0]    
    bs = False
    linear_fit = True
    quadratic_fit = True
#print dx
#print dy
corr = None
pars_lin = None 
pars_quad = None


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

_stream = None

instances = []

def _get_device(d):
    global _stream
    egu = None
    if isinstance(d, basestring):    
        name = d.strip()
        d = None
        try:
            d = get_device(name) 
            if d is None: 
                d = eval(name)
            #print name
            if d is not None:
                if not isinstance(r, Device): 
                    d = None
                else:
                    try:
                        egu = d.unit
                    except:
                        pass                            
        except:
            pass
        if d is None:
            offset = 0
            if " " in name:
                tokens = name.split(" ")
                name = tokens[0]
                offset = int(tokens[1])
            if bs == True:
                if _stream == None:
                    _stream = Stream("corr_stream", dispatcher)
                    instances.append(_stream)
                d = _stream.addScalar(name, name, int(interval*100), offset)
            else:
                d =  Channel(name)    
                instances.append(d)
            try:
                egu = caget(name+".EGU",'s')
            except:
                pass            
    else:
        try:
            egu = d.unit
        except:
            pass            
        
    return d, egu

dx, egux = _get_device(dx)
dy, eguy = _get_device(dy)

p.getAxis(p.AxisId.X).setLabel(egux)
p.getAxis(p.AxisId.Y).setLabel(eguy)


try:
    if _stream != None:
        _stream.initialize()
        _stream.start(True)
    p.addSeries(LinePlotSeries("Data")) 
    sd=p.getSeries(0)
    sd.setLinesVisible(False)    
    
    sd.setPointSize(4)

    if get_exec_pars().source == CommandSource.ui:
        if globals().has_key("marker"):
            p.removeMarker(marker)    
        marker=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 get_exec_pars().source == CommandSource.ui:
                    if marker is not None: 
                        p.removeMarker(marker)
                    marker = p.addMarker(x2+res, p.AxisId.X, s,  p.getBackground())
                    marker.setLabelPaint(getColorStdout())
                if linear_fit:
                    #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")
                if quadratic_fit:
                    #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()
        stop_task("outupdate")