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.32
    energy0 = 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]
    energy0 = args[6]
    plt     = args[7]

A = energy0 / disp / 1e6
B = energy0

#Plot setup
plt.clear()
plt.setStyle(plt.Style.ErrorY)
plt.addSeries(LinePlotErrorSeries("Energy"))
plt.addSeries(LinePlotErrorSeries("Energy Spread", None, 2))
plt.getAxis(plt.AxisId.X).setLabel("Gun Beam Phase (deg)")
plt.getAxis(plt.AxisId.Y).setLabel("Energy (MeV)")
plt.getAxis(plt.AxisId.Y2).setLabel("Energy Spread (MeV)")
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  = "SLG-LCAM-C041"
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_center_of_mass")
dx  = camtool.stream.getChild("x_rms")

#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  
    E_mean,  E_stdev  = A * x_pos_mean + B,    abs(A) * x_pos_stdev
    dE_mean, dE_stdev = abs(A) * x_width_mean, abs(A) * x_width_stdev
    plt.getSeries(0).appendData(record.positions[0], E_mean, E_stdev)
    plt.getSeries(1).appendData(record.positions[0], dE_mean, dE_stdev)

#The scan loop
try:
    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
E  = [A * val.mean + B  for val in r.getReadable(0)]
dE = [abs(A) * val.mean for val in r.getReadable(1)]
save_dataset(get_exec_pars().group + "/E",  E)
save_dataset(get_exec_pars().group + "/dE",  dE)

if do_elog:
    if get_option("Generated data file:\n" + get_exec_pars().path +"\n\n" + "Save to ELOG?", "YesNo") == "Yes":    
        log_msg = "Data file: " + get_exec_pars().path      
        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() + "/GunScanPlot.png")
        plt.saveSnapshot(file_name , "png")        
        elog("Gun scan", log_msg, [file_name,])