import  ch.psi.pshell.epics.ControlledVariable as ControlledVariable

if get_exec_pars().source == CommandSource.ui:
    station = "SINSB04"
    bpm_ch  = "SINBC01-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 =  ControlledVariable("Phase", station + "-RSYS:SET-VSUM-PHASE", station + "-RSYS:GET-VSUM-PHASE")
phase.config.minValue   =-45.0
phase.config.maxValue   = 360.0
phase.config.resolution = 0.5
phase.initialize()

V =  Channel(station + "-RSYS:GET-VSUM-AMPLT",     type = 'd', alias='Amplitude Readback')
P =  Channel(station + "-RSYS:GET-KLY-POWER",      type = 'd', alias='Power Readback')
x =  Channel(bpm_ch  + ":X1",                      type = 'd', alias='BPM-X')

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'))
caput(station + "-RSYS:CALC-VSUM-PHASE-OFFSET", float('nan'))
caput(station + "-RSYS:CALC-VSUM-AMPLT-SCALE" , float('nan'))
caput(station + "-RSYS:CALC-VOLT-POWER-SCALE" , 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(rec.values[1].mean/1000.0/disp*energy0)
   caput(station + "-RSYS:GET-ENERGY-ARRAY", to_array(arr_energy,'d'))
   caput(station + "-RSYS:GET-PHASE-ARRAY",  to_array(arr_phase, 'd'))

try:
    Vb = create_averager(V, nb, 0.100)
    xb = create_averager(x, nb, 0.100)
    r = lscan(phase, [Vb, xb], start, stop, step , latency=lat, after_read = after)
    rf_phase = r.getPositions(0)
    energy   = [val.mean/1000.0/disp*energy0 for val in r.getReadable(1)]
    caput(station + "-RSYS:GET-ENERGY-ARRAY", to_array(energy,  'd'))
    caput(station + "-RSYS:GET-PHASE-ARRAY",  to_array(rf_phase,'d'))
    phase_fit_max = None
    try:
        (energy_max, angular_frequency, phase0, in_range, phase_fit_max, fit_x, fit_y) = hfit(energy , xdata = rf_phase)
    except:
        raise Exception("Fit failure")
    caput(station + "-RSYS:GET-ONCREST-VSUM-PHASE", phase_fit_max)
    caput(station + "-RSYS:GET-ONCREST-E-GAIN",     energy_max)
    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 <= phase_fit_max <= phase_max):
        raise Exception("Fit maximum outside scan range")
    phase.write(phase_fit_max)
    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.close()
    V.close()
    P.close()
    x.close()

print ("------------------------------------")
print ("Valid fit")

energy_gain     = energy_max
phase_offset    = 90 - phase_fit_max
amplitude_scale = energy_gain / Ampl
power_scale     = Power / math.pow(Ampl,2)

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

#title="Phase scan "+str(station)
#message=("Energy Gain:     %0.3f"   % energy_gain     + "MeV\n" +
#         "Phase Offset:    %0.2f"   % phase_offset    + "deg\n" +
#         "Amplitude Scale: %0.3f"   % amplitude_scale + "MV\n" +
#         "Power Scale:     %0.3f"   % power_scale)    + "1/ohm"
#elog(title, message)