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
import ch.psi.pshell.bs.StreamMerger as StreamMerger


_stop_exec = False  
_task = None
_dispatcher_stream = None

def start_camera_correlation(dispatcer_channel, camera_channel, samples=50, modulo=1, offset=0, plt=None):
    global _stop_exec, _task, _dispatcher_stream
    stop_camera_correlation()

    if plt:
        for s in plt.getAllSeries():
            plt.removeSeries(s)    
    
    class LinearFit(ReadonlyAsyncRegisterBase, ReadonlyRegisterArray):
        def getSize(self):        
            return 2        
        def set(self, pars):
            self.onReadout(to_array(pars, 'd'))
    
    class Correlation(ReadonlyAsyncRegisterBase):
        def set(self, val):
            self.onReadout(val)
    
    add_device(LinearFit("linear_fit"), True)
    add_device(Correlation("correlation"), True)
    
    
    
    try:
        _dispatcher_stream = Stream("corr_stream", dispatcher)
        _dispatcher_stream.addScalar(dispatcer_channel, dispatcer_channel, modulo, offset)    
        _dispatcher_stream.initialize()
        _dispatcher_stream.start(True)
        _dispatcher_stream.setBufferCapacity(500)
    
        camera_name, camera_channel = camera_channel.split(" ")
        shared = camera_name.endswith("_sp1")
        cam_server.start(camera_name, shared )            
        camera_stream = cam_server.stream
        camera_stream.setBufferCapacity(500)
        camera_stream.waitCacheChange(10000);
    
        def merger_task():  
            merger = StreamMerger("stream", _dispatcher_stream, camera_stream)
            try:
                if plt:
                    sd = LinePlotSeries("Data")
                    plt.addSeries(sd) 
                    sd.setLinesVisible(False)        
                    sd.setPointSize(4)
                else:
                    ax,ay = [],[]        
                merger.monitored=True
                merger.start()           
                merger.waitCacheChange(5000)
                
                dx=merger.getChild(dispatcer_channel)
                dy=merger.getChild(camera_channel)    
            
                while(not _stop_exec):
                    merger.waitCacheChange(10000)
                    x=dx.read() 
                    y=dy.read()
        
                    if plt:
                        sd.appendData(x, y)
                        while len(sd.x) > samples:            
                            sd.token.remove(0)  #Remove First Element  
                        ax,ay = sd.x,sd.y
                    else:
                        ax.append(x); ay.append(y)
                        while len(ax) > samples:            
                            ax.pop(0); ay.pop(0)               
                               
                    if len(ax)>2:
                        x1, x2 = min(ax), max(ax)  
                        if x1!=x2:
                            #Display correlation
                            corr= PearsonsCorrelation().correlation(to_array(ax,'d'), to_array(ay,'d'))        
                            correlation.set(corr)
                            pars_lin = (a0,a1) = fit_polynomial(ay, ax, 1)
                            linear_fit.set(pars_lin)
                            y1,y2 = poly(x1, pars_lin), poly(x2, pars_lin)
                            if plt:
                                invoke((plot_line,(plt, x1, y1, x2, y2, 2, Color.BLUE, "Fit Linear")), False)
            finally:
                merger.close()
        _task = fork(merger_task)
    except:
        stop_camera_correlation()
        raise

def stop_camera_correlation():
    global _stop_exec, _task, _dispatcher_stream
    _stop_exec = True
    if _task:
       join(_task) 
    cam_server.stop()
    if _dispatcher_stream:
        _dispatcher_stream.close()
    _stop_exec = False  
    _task = None
    _dispatcher_stream = None        


#Testing
start_camera_correlation("SINEG01-DICT215:B1_CHARGE", "SATES31-CAMS187-RIXS1_sp1 intensity", samples = 40, plt=plot(None)[0])
try:
   time.sleep(10.0)
finally:
   stop_camera_correlation()