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

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:
    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]
else:
    start   = 44.0
    stop    = 65.0
    step    = 1.0
    nb      = 3
    lat     = 0.3
    disp    = -0.387
    p0      = 7.1
    plt     = plot(None, title="Output")[0]

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("Beam phase", "SINEG01-RSYS:SET-BEAM-PHASE-SIM", "SINEG01-RSYS:SET-BEAM-PHASE-SIM")
    camera_name  = "simulation"
else:
    phase = Positioner("Beam phase", "SINEG01-RSYS:SET-BEAM-PHASE", "SINEG01-RSYS:GET-BEAM-PHASE")
    camera_name  = "SINBD01-DSCR010"
phase.config.minValue   = -360.0
phase.config.maxValue   = 360.0
phase.config.precision  = 3
phase.config.resolution = 0.1
phase.config.rotation   = False
phase.config.save()
phase.initialize()
phase0 = phase.read()

#Camera setup
cam_server.start(camera_name)
wait_cam_server_message()
x   = cam_server.stream.getChild("x_fit_mean")
dx  = cam_server.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, plt
    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()
    cam_server.stop()   # stops cam_server but does not close it cam_server is a global object

ph = r.getPositions(0)
p  = [A * val.mean + B  for val in r.getReadable(0)]
dp = [abs(A) * val.mean for val in r.getReadable(1)]

#Fitting and plotting
try:
    i_max = p.index(max(p))
    i_min = dp.index(min(dp))
    min_i, max_i = max(i_max-5, 0), min(i_max+6, len(p))
    (ph_p_max, p_max, ph_p_fit, p_fit)  = extremum(ph[min_i:max_i], p[min_i:max_i])
    min_i, max_i = max(i_min-5, 0), min(i_min+6, len(dp))
    (ph_dp_min, dp_min, ph_dp_fit, dp_fit) = extremum(ph[min_i:max_i], dp[min_i:max_i])
    plt.addSeries(LinePlotErrorSeries("Momentum Fit", plt.getSeries(0).color))
    plt.addSeries(LinePlotErrorSeries("Momentum Spread Fit", plt.getSeries(1).color, 2))
    plt.getSeries(2).setData(ph_p_fit, p_fit)
    plt.getSeries(3).setData(ph_dp_fit, dp_fit)    
    plt.getSeries(2).setPointsVisible(False)
    plt.getSeries(3).setPointsVisible(False)
    plt.addMarker(ph_p_max, plt.AxisId.X, "%3.2f" % ph_p_max, plt.getSeries(0).color)
    plt.addMarker(ph_dp_min, plt.AxisId.X, "%3.2f" % ph_dp_min, plt.getSeries(1).color)
except:
    raise Exception("Fit failure")

#Saving metadata
save_dataset(get_exec_pars().group + "/p",  p)
set_attribute(get_exec_pars().group + "/p", "ph_p_max", ph_p_max)
set_attribute(get_exec_pars().group + "/p", "p_max", p_max)
save_dataset(get_exec_pars().group + "/dp",  dp)
set_attribute(get_exec_pars().group + "/dp", "ph_dp_min", ph_dp_min)
set_attribute(get_exec_pars().group + "/dp", "dp_min", dp_min)

#Elog entry
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 at %0.2f" % phase0 + " deg beam phase\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,])