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

if get_exec_pars().source == CommandSource.ui:
    prefix = "SINSB04-RSYS"
else:
    prefix = args[0] + "-RSYS"

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

phase =  ControlledVariable("Phase", prefix + ":SET-VSUM-PHASE-SIM", prefix + ":GET-VSUM-PHASE-SIM")
phase.config.minValue   =-45.0
phase.config.maxValue   = 360.0
phase.config.resolution = 0.5
phase.initialize()

V =  Channel(prefix + ":GET-VSUM-AMPLT-SIM", type = 'd', alias='Amplitude Readback')
P =  Channel(prefix + ":GET-KLY-POWER-SIM",  type = 'd', alias='Power Readback')

x =  Channel(bpm_ch + ":X1-SIMU",            type = 'd', alias='BPM-X')

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

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

try:
    Vb = create_averager(V, nb, lat)
    xb = create_averager(x, nb, lat)

    r = lscan(phase, [V, x], start, stop, step , latency=lat, after_read = after)
    rf_phase = r.getPositions(0)
    energy   = [x/1000.0/disp*energy0 for x in r.getReadable(1)]
    caput(prefix + ":GET-ENERGY-ARRAY", to_array(energy,'d'))
    caput(prefix + ":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(prefix + ":GET-ONCREST-VSUM-PHASE", phase_fit_max)
    caput(prefix + ":GET-ONCREST-E-GAIN",     energy_max)
    caput(prefix + ":GET-FIT-PHASE-ARRAY",    fit_x)
    caput(prefix + ":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)
    V = V.read()
    P = P.read()
    caput(prefix + ":GET-ONCREST-VSUM-AMPLT", V)
    caput(prefix + ":GET-ONCREST-KLY-POWER",  P)
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 / V
power_scale     = P / math.pow(V,2)

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

set_return("\nEnergy Gain: "     + str(energy_gain)     + "\n" +
           "Phase Offset: "      + str(phase_offset)    + "\n" +
           "Amplitude Scale: "   + str(amplitude_scale) + "\n" +
           "Power Scale: "       + str(power_scale))