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

#Parameters
dry_run = False
do_elog = True
if get_exec_pars().source == CommandSource.ui:
    start   = 45.0
    stop    = 55.0
    step    = 1.0
    nb      = 3
    lat     = 0.3
    disp    = -0.387
    p0      = 7.1
    plt     = plot(None, title="Output")[0]
else:
    start   = args[0]
    stop    = args[1]
    step    = args[2]
    nb      = int(args[3])
    lat     = args[4]
    disp    = args[5]
    p0      = args[6]
    plt     = args[7]

A = p0 / disp / 1e6
B = p0

#Plot setup
plt.clear()
plt.setStyle(plt.Style.ErrorY)
plt.addSeries(LinePlotErrorSeries("Momentum"))
plt.addSeries(LinePlotErrorSeries("Momentum Spread", None, 2))
plt.getAxis(plt.AxisId.X).setLabel("Gun Beam Phase (deg)")
plt.getAxis(plt.AxisId.Y).setLabel("Momentum (MeV/c)")
plt.getAxis(plt.AxisId.Y2).setLabel("Momentum Spread (MeV/c)")
plt.setLegendVisible(True)

#Creating Phase positioner
if dry_run:
    phase = Positioner("Phase", "SINEG01-RSYS:SET-BEAM-PHASE-SIM", "SINEG01-RSYS:SET-BEAM-PHASE-SIM")
    camera_name  = "simulation"
else:
    phase = Positioner("Phase", "SINEG01-RSYS:SET-BEAM-PHASE", "SINEG01-RSYS:GET-BEAM-PHASE")
    camera_name  = "SINBD01-DSCR010"

phase.config.minValue   = -180.0
phase.config.maxValue   = 180.0
phase.config.precision  = 3
phase.config.resolution = 0.01
phase.config.rotation   = False
phase.config.save()
phase.initialize()
phase0 = phase.read()

#Camtool setup
kill_camtool()
check_camtool()
camtool.start(camera_name)
wait_camtool_message()
x   = camtool.stream.getChild("x_fit_mean")
dx  = camtool.stream.getChild("x_fit_standard_deviation")

#Creating averagers
x_averager  = create_averager(x,  nb, -1)    # -1 event based, waits for the next value
dx_averager = create_averager(dx, nb, -1)
dx_averager.monitored=True     # not blocking, will return last nb values

#Record callback: uptate of output plot
def after_sample(record, scan):
    global A, B
    x_pos_mean,   x_pos_stdev   = record.values[0].mean, record.values[0].stdev
    x_width_mean, x_width_stdev = record.values[1].mean, record.values[1].stdev
    p_mean,  p_stdev  = A * x_pos_mean + B,    abs(A) * x_pos_stdev
    dp_mean, dp_stdev = abs(A) * x_width_mean, abs(A) * x_width_stdev
    plt.getSeries(0).appendData(record.positions[0], p_mean, p_stdev)
    plt.getSeries(1).appendData(record.positions[0], dp_mean, dp_stdev)

#The scan loop
try:
    phase.write(start)
    time.sleep(1.0)
    r = lscan(phase, [x_averager, dx_averager], start, stop, step , latency=lat, after_read = after_sample)
finally:
    phase.write(phase0)
    phase.close()
    camtool.stop()   # stops camtool but does not close it camtool is a global object

#Saving metadata
p  = [A * val.mean + B  for val in r.getReadable(0)]
dp = [abs(A) * val.mean for val in r.getReadable(1)]
save_dataset(get_exec_pars().group + "/p",  p)
save_dataset(get_exec_pars().group + "/dp",  dp)

if do_elog:
    if get_option("Generated data file:\n" + get_exec_pars().path +"\n\n" + "Save to ELOG?", "YesNo") == "Yes":
        Laser = str(caget("SLGTV-LMTO-M055:MOT-KNOWN-POS"))
        log_msg = "Data file: " + get_exec_pars().path + "\n\n"
        log_msg = log_msg + "Laser:            " + Laser + "\n"
        if Laser == "Alcor":
            log_msg = log_msg + "Energy plate:     %0.2f" % caget("SLGTH01-LMRM-M074:MOT.RBV") + " deg \n"
        else:
            log_msg = log_msg + "Energy plate:     %0.2f" % caget("SLGJG-LMRM-M031:MOT.RBV") + " deg \n"
        if caget("SLGTV-LMTO-M053:MOT-ACT-POS") == "IRIS":
            log_msg = log_msg + "Collimator:       IRIS %0.2f" % caget("SLGTV-LAPP:SIZE-GET") + " mm \n"
        else:
            log_msg = log_msg + "Collimator:       " + str(caget("SLGTV-LMTO-M053:MOT-ACT-POS")) + "\n"
        log_msg = log_msg + "Charge:           %0.2f" % caget("SINEG01-DICT215:B1_CHARGE-OP") + " pC \n"
        sleep(0.1) #Give some time to plot to be finished - it is not sync  with acquisition
        file_name = os.path.abspath(get_context().setup.getContextPath() + "/GunEnergyScanPlot.png")
        plt.saveSnapshot(file_name , "png")
        elog("Gun Energy Scan", log_msg, [file_name,])