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

dry_run = True
do_elog = False

if get_exec_pars().source == CommandSource.ui:
    station = "STEST01"
    bpm_ch  = "SINBC02-DBPM140"
else:
    station = args[0]
    bpm_ch  = args[1]

start   = caget(station + "-RSYS:SET-SCAN-START")
stop    = caget(station + "-RSYS:SET-SCAN-STOP")
step    = caget(station + "-RSYS:SET-SCAN-STEP")
lat     = caget(station + "-RSYS:SET-SCAN-WAIT-TIME")
nb      = caget(station + "-RSYS:SET-NUM-AVERAGE")
disp    = caget(bpm_ch  + ":DISPERSION")
energy0 = caget(bpm_ch  + ":ENERGY")

phase =  Positioner("Phase", station + "-RSYS:SET-VSUM-PHASE", station + "-RSYS:GET-VSUM-PHASE")
phase.config.minValue   =-180.0
phase.config.maxValue   = 180.0
phase.config.precision  = 3
phase.config.rotation   = True
phase.config.resolution = 0.5
phase.initialize()
V =  ChannelDouble("Amplitude Readback", station + "-RSYS:GET-VSUM-AMPLT")
P =  ChannelDouble("Power Readback",     station + "-RSYS:GET-KLY-POWER")
if dry_run:
    x =  ChannelDouble("BPM-X", bpm_ch  + ":X1-SIMU")
else:
    x =  ChannelDouble("BPM-X", bpm_ch  + ":X1")
V.initialize()
P.initialize()
x.initialize()

phase0 = phase.read()

caput(station + "-RSYS:GET-FIT-PHASE-ARRAY",    to_array([0.0],'d'))
caput(station + "-RSYS:GET-FIT-ENERGY-ARRAY",   to_array([0.0],'d'))
caput(station + "-RSYS:GET-ONCREST-VSUM-PHASE", float('nan'))
caput(station + "-RSYS:GET-ONCREST-VSUM-AMPLT", float('nan'))
caput(station + "-RSYS:GET-ONCREST-E-GAIN",     float('nan'))
caput(station + "-RSYS:GET-ONCREST-KLY-POWER",  float('nan'))

#update the plot dynamically
arr_phase,arr_energy = [],[]
def after(rec):
   global disp, energy0
   arr_phase.append(rec.positions[0])
   arr_energy.append(energy0 * (1 + rec.values[0].mean / 1000.0 / disp))
   caput(station + "-RSYS:GET-PHASE-ARRAY",  to_array(arr_phase, 'd'))
   caput(station + "-RSYS:GET-ENERGY-ARRAY", to_array(arr_energy,'d'))

try:
    x_averager = create_averager(x, nb, 0.100)
    r = lscan(phase, x_averager, start, stop, step, latency=lat, after_read = after)
    rf_phase = r.getPositions(0)
    energy   = [energy0 * (1 + val.mean / 1000.0 / disp) for val in r.getReadable(0)]
    caput(station + "-RSYS:GET-ENERGY-ARRAY", to_array(energy,  'd'))
    caput(station + "-RSYS:GET-PHASE-ARRAY",  to_array(rf_phase,'d'))
    try:
        run("CPython/wrapper")
        (fit_amplitude, fit_phase_deg, fit_offset, ph_crest, fit_x, fit_y) = hfitoff(energy , rf_phase)
    except:
        raise Exception("Fit failure")
    plot([energy, fit_y], ["data", "fit"], [rf_phase, fit_x])
    caput(station + "-RSYS:GET-ONCREST-VSUM-PHASE", ph_crest)
    caput(station + "-RSYS:GET-ONCREST-E-GAIN",     fit_amplitude)
    caput(station + "-RSYS:GET-FIT-PHASE-ARRAY",    fit_x)
    caput(station + "-RSYS:GET-FIT-ENERGY-ARRAY",   fit_y)
    phase_min, phase_max = min(rf_phase), max(rf_phase)
    if not (phase_min <= ph_crest <= phase_max):
        raise Exception("On-crest phase outside scan range")
    phase.write(ph_crest)
    time.sleep(lat)
    Ampl  = V.read()
    Power = P.read()
    caput(station + "-RSYS:GET-ONCREST-VSUM-AMPLT", Ampl)
    caput(station + "-RSYS:GET-ONCREST-KLY-POWER",  Power)
finally:
    phase.write(phase0)
    phase.close()
    V.close()
    P.close()
    x.close()

phase_offset    = 90 - ph_crest
amplitude_scale = fit_amplitude / Ampl
power_scale     = Power / math.pow(Ampl,2)

caput(station + "-RSYS:SET-VSUM-PHASE-OFFSET-BASE",  phase_offset)
caput(station + "-RSYS:SET-VSUM-AMPLT-SCALE",        amplitude_scale)
caput(station + "-RSYS:SET-VOLT-POWER-SCALE",        power_scale)

if do_elog:
    if get_option("Generated data file:\n" + get_exec_pars().path +"\n\n" + "Save to ELOG?", "YesNo") == "Yes":
        title = "Phase scan" + station
        log_msg =           "Data file:               " + get_exec_pars().path      + "\n"
        log_msg = log_msg + "Energy Gain:     %0.3f"   % energy_gain     + "MeV\n"
        log_msg = log_msg + "Phase Offset:    %0.2f"   % phase_offset    + "deg\n"
        log_msg = log_msg + "Amplitude Scale: %0.3f"   % amplitude_scale + "MV\n"
        log_msg = log_msg + "Power Scale:     %0.3f"   % power_scale     + "MW/MV^2"
        attachments = []
        if plt is not None:
            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() + "/SchottkyScanPlot.png")
            plt.saveSnapshot(file_name, "png")
            attachments = [file_name,]
        elog(title, log_msg, attachments)