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

dry_run = False
do_elog = True

if get_exec_pars().source == CommandSource.ui:
    station = "SINSB01"
    bpm_ch  = "S10BD01-DBPM020"
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")

A = energy0 / disp / 1e3
B = energy0

phase =  Positioner("Phase", station + "-RSYS:SET-VSUM-PHASE", station + "-RSYS:GET-VSUM-PHASE")
phase.config.minValue   =-360.0
phase.config.maxValue   = 360.0
phase.config.precision  = 4
phase.config.rotation   = False
phase.config.resolution = 0.1
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 A, B
   arr_phase.append(rec.positions[0])
   arr_energy.append(A * rec.values[0].mean + B)
   caput(station + "-RSYS:GET-PHASE-ARRAY",  to_array(arr_phase, 'd'))
   caput(station + "-RSYS:GET-ENERGY-ARRAY", to_array(arr_energy,'d'))

try:
    phase.write(start)
    time.sleep(1.0)
    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   = [A * val.mean + B 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")
    """
    # This only works when pshell is visible not in server mode to be fixed
    plt = plot(None,name="data")[0]    
    plt.getSeries(0).setData(to_array(rf_phase,'d'), to_array(energy,'d'))
    plt.getSeries(0).setPointSize(8)    
    plt.getSeries(0).setLinesVisible(False)
    plt.addSeries(LinePlotSeries("fit"))
    plt.getSeries(1).setData(fit_x, fit_y)
    plt.getSeries(1).setPointsVisible(False)
    plt.setLegendVisible(True)
    """
    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.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.0 - ph_crest
amplitude_scale = fit_amplitude / Ampl
power_scale     = Power / math.pow(fit_amplitude,2)

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

if do_elog:
    title = "Phase scan" + station
    log_msg =           "Data file:          " + get_exec_pars().path   + "\n"
    log_msg = log_msg + "On-crest VS phase:  %0.2f"   % ph_crest        + " deg \n"
    log_msg = log_msg + "Energy Gain:        %0.3f"   % fit_amplitude   + " MeV \n"
    log_msg = log_msg + "VS -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 = get_plot_snapshots()
    elog(title, log_msg, attachments)