import ch.psi.pshell.epics.Positioner as Positioner

dry_run = False
do_elog = True
is_panel = get_exec_pars().source != CommandSource.ui    # must be check before run

run("CPython/wrapper")
#Parameters
if is_panel:
    MCR1    = args[0]
else:
    U       = "X"
    MQUA    = "SATMA01-MQUA050"
    MCR1    = {"X": "SATCB01-MCRX430",        "Y": "SATCB01-MCRY430"       }
    BPM1    = {"X": "SARCL01-DBPM150:X1",     "Y": "SARCL01-DBPM150:Y1"    }
    BPM2    = {"X": "SARCL02-DBPM110:X1",     "Y": "SARCL02-DBPM110:Y1"    }
    dC      = 2.0
    dK      = 0.05
    step    = 10
    nb      = 5
    lat     = 0.0
    plt     = plot(None, title="Output")[0]
    
#Plot setup
plt.clear()
plt.removeMarker(None)
plt.setStyle(plt.Style.ErrorXY)
plt.addSeries(LinePlotErrorSeries("run0", Color.red))
plt.addSeries(LinePlotErrorSeries("run1", Color.yellow))
plt.addSeries(LinePlotErrorSeries("run2", Color.green))
plt.getAxis(plt.AxisId.X).setLabel("BPM1 (mm)")
plt.getAxis(plt.AxisId.Y).setLabel("BPM2 (mm)")
plt.setLegendVisible(True)
    
#Creating Phase positioner
if dry_run:
    kick1 = Positioner("kick1", "SATCB01-MCRX430:I-SET", "SATCB01-MCRX430:I-READ")
    quad = Positioner("quad", "SATMA01-MQUA050:K1L-SET", "SATMA01-MQUA050:K1L-READ")
else:
    kick1 = Positioner("kick1", MCR1[U] + ":I-SET", MCR1[U] + ":I-READ")
    quad = Positioner("quad", MQUA + ":K1L-SET", MQUA + ":K1L-READ")
quad.config.minValue   = -2.0
quad.config.maxValue   = 2.0
kick1.initialize()
quad.initialize()
kick1_ini = kick1.read()
quad_ini = quad.read()
start = kick1_ini - dC / 2
stop = kick1_ini + dC / 2

#Creating averagers
X1  = "ca://" + BPM1[U]
X2  = "ca://" + BPM2[U]
X1_averager = create_averager(X1,  nb, -1)    # -1 event based, waits for the next value
X2_averager = create_averager(X2, nb, -1)     # ca:// it comes from epics channels
X2_averager.monitored = True                  # not blocking, will return last nb values

#Record callback: uptate of output plot
def after_sample(record, scan):
    global plt,i
    X1_mean, X1_stdev = record[0].mean, record[0].stdev
    X2_mean, X2_stdev = record[1].mean, record[1].stdev
    plt.getSeries(i).appendData(X1_mean, X2_mean, X1_stdev, X2_stdev)
    
#The scan loop
try:
    data = []
    i = 0
    for K1L in (quad_ini, quad_ini - dK, quad_ini + dK):
        kick1.write(start)
        quad.write(K1L)
        time.sleep(2.0)
        r = lscan(kick1, [X1_averager, X2_averager], start, stop, step , latency=lat, after_read = after_sample)
        data.append({"K1L":K1L,"scan":r})
        i = i + 1
finally:
    kick1.write(kick1_ini)
    kick1.close()
    quad.write(quad_ini)
    quad.close()

#Fitting and plotting
P = []
for i in range(len(data)):
    x = data[i]["scan"].getReadable(0)
    y = data[i]["scan"].getReadable(1)
    (p, x_fit, y_fit, R2) = linfit(x, y)
    P.append(p)
    plt.addSeries(LinePlotErrorSeries("Fit"+str(i), plt.getSeries(i).color))
    plt.getSeries(i+3).setData(x_fit, y_fit)
    plt.getSeries(i+3).setPointsVisible(False)
a, b = P[1][0], P[1][1] #second line +K
c, d = P[2][0], P[2][1] #third line -K
ui = (d - b)/(a - c) if a != c else None
    
#Elog entry
if do_elog:
    if get_option("Generated data file:\n" + get_exec_pars().path +"\n\n" + "Save to ELOG?", "YesNo") == "Yes":
        title = "BPM_Quad_scan"
        log_msg = "Data file: " + get_exec_pars().path + "\n\n"    
        log_msg = log_msg + "MQUA:                  " + MQUA + "\n"
        log_msg = log_msg + "BPM1:                  " + BPM1[U] + "\n"
        log_msg = log_msg + "BPM2:                  " + BPM2[U] + "\n"
        log_msg = log_msg + "MQUA center:           " + str(ui) + "\n"
        sleep(0.1) #Give some time to plot to be finished - it is not sync  with acquisition
        file_name = os.path.abspath(Setup.getContextPath() + "/BPM_Quad_Plot.png")
        plt.saveSnapshot(file_name , "png")
        elog(title, log_msg, [file_name,])

set_return(ui)