"""
Arguments:
DRY_RUN (BOOL)
CENTER_X, CENTER_Y (FLOAT)
STEPS_X, STEPS_Y  (INT)
STEP_SIZE_X, STEP_SIZE_Y (FLOAT)
EXPOSURES (list or None)
SETTLING_TIME (double)
ZIGZAG (BOOLEAN)
COMPRESSION (BOOLEAN)
NOISE (FLOAT)
"""
import random

position_plot = None if "POSITION_PLOT" not in globals() else POSITION_PLOT

#Debugging
if (get_exec_pars().args is None) and (get_exec_pars().script=="Regine"):
    DRY_RUN = True
    CENTER_X = 0.0
    CENTER_Y = 0.0
    EXPOSURES =  [0.1, 0.5, 1.0]
    STEPS_X = 5
    STEPS_Y= 3
    STEP_SIZE_X = 0.1
    STEP_SIZE_Y = 0.2
    SETTLING_TIME = 1.0
    ZIGZAG = True
    COMPRESSION = False
    NOISE=0.10
    position_plot=None

           
if not position_plot: position_plot=plot(None, title="Motor Positions")[0]

current=sin
ccd=scienta
dummy_x, dummy_y = motor, motor2
sample_x, sample_y = motor, motor2
after_readout= before_readout=None 
def after_scan():
    pass

#Scanning exposure
class exposure_index (Writable, Readable):
    def __init__(self):            
        self.pos=0.0
                  
    def read(self):
        return self.pos

    def write(self, value): 
        ccd.setExposure(EXPOSURES[int(value)])
        self.pos = value
        
class exposure (Readable):            
    def read(self):
        return ccd.exposure


#Constants
SENSORS = [current, ccd.dataMatrix]
POSITIONERS =  [dummy_x, dummy_y] if DRY_RUN else [sample_x, sample_y]

ENABLED_PLOTS = [current, ccd.dataMatrix]
CUSTOM_PLOT_TYPES = {ccd.dataMatrix:"ch.psi.pshell.plot.MatrixPlotRenderer"}

RANGE_X=[CENTER_X - STEP_SIZE_X*STEPS_X,  CENTER_X + STEP_SIZE_X*STEPS_X]
RANGE_Y=[CENTER_Y - STEP_SIZE_Y*STEPS_Y,  CENTER_Y + STEP_SIZE_Y*STEPS_Y]


position_plot.clear()
position_plot.addSeries(LinePlotSeries("positions"))
position_plot.getAxis(AxisId.X).label = POSITIONERS[0].name
position_plot.getAxis(AxisId.X).setRange(RANGE_X[0]-STEP_SIZE_X, RANGE_X[1]+STEP_SIZE_X)
position_plot.getAxis(AxisId.Y).label = POSITIONERS[1].name
position_plot.getAxis(AxisId.Y).setRange(RANGE_Y[0]-STEP_SIZE_Y, RANGE_Y[1]+STEP_SIZE_Y)


def after_read(record, scan):
    position_plot.getSeries(0).appendData(record[POSITIONERS[0]], record[POSITIONERS[1]])


if EXPOSURES:
    #ENABLED_PLOTS = [ccd.dataMatrix] #Cannot plot current in 3d vscan
    POSITIONERS = POSITIONERS + [exposure_index()]
    SENSORS = SENSORS  + [exposure()]
    RANGE_E=[0, len(EXPOSURES)-1]

    #set_exec_pars(manual_range=RANGE_E, manual_range_y=RANGE_Y)

    CUSTOM_PLOT_TYPES[sin]=1
    set_exec_pars(manual_range=RANGE_X)
    #set_exec_pars(domain_axis="Index")    
    #set_exec_pars(auto_range=True)
else:
    set_exec_pars(manual_range=RANGE_X, manual_range_y=RANGE_Y)

def gen():
    x_index = y_index = e_index = 0
    for x_step in range(-STEPS_X, STEPS_X+1):
        xpos = CENTER_X + x_step*STEP_SIZE_X + (random.random()-0.5)*NOISE*2*STEP_SIZE_X
        range_y=  range(STEPS_Y, -STEPS_Y-1, -1) if (ZIGZAG and(x_index%2 ==1)) else range(-STEPS_Y, STEPS_Y+1)
        for y_step in range_y:            
            ypos = CENTER_Y + y_step*STEP_SIZE_Y + (random.random()-0.5)*NOISE*2*STEP_SIZE_Y        
            if EXPOSURES:
                range_e=  range(len(EXPOSURES)-1,-1, -1) if (ZIGZAG and(y_index%2 ==1)) else range(len(EXPOSURES))
                for e in range_e:
                    yield([xpos,ypos,e])
            else:    
                yield([xpos,ypos])
            y_index=y_index+1
        x_index=x_index+1


    
try:
    r = vscan(POSITIONERS, SENSORS , gen(), False,\ 
          SETTLING_TIME, relative=False, zigzag = ZIGZAG, initial_move=False, \
          before_read=before_readout, after_read = after_readout, \
          manual_range=RANGE_X, manual_range_y=RANGE_Y, \
          compression = COMPRESSION, enabled_plots=ENABLED_PLOTS, \
          keep=False, check_positions=False, plot_types=CUSTOM_PLOT_TYPES)  
    set_return(r)  
finally:
    after_scan()