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

if get_context().source != CommandSource.ui:
    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")

rf_phase_var =  ControlledVariable("Phase", prefix + ":SET-VSUM-PHASE-SIM", prefix + ":SET-VSUM-PHASE-SIM") #TODO: CHANGE TO GET
rf_phase_var.config.minValue   =-180.0
rf_phase_var.config.maxValue   = 180.0
rf_phase_var.config.resolution = 0.001
rf_phase_var.initialize()

rf_ampl_rbk  =  Channel(prefix + ":GET-VSUM-AMPLT-SIM", type = 'd', alias='Amplitude Readback')
rf_ampl_rbk.write(1.0)
rf_power_rbk =  Channel(prefix + ":GET-KLY-POWER-SIM",  type = 'd', alias='Power Readback')
bpm_x        =  Channel(bpm_ch + ":X1-SIM",                 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 + ":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:
    r = lscan(rf_phase_var, [rf_ampl_rbk, bpm_q], 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-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")

    rf_phase_var.write(phase_fit_max)
    time.sleep(lat)
    ampl  = rf_ampl_rbk.read()
    power = rf_power_rbk.read()
finally:
    rf_phase_var.close()
    rf_ampl_rbk.close()
    rf_power_rbk.close()
    bpm_x.close()

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

energy_gain     = energy_max
phase_offset    = - phase_fit_max
amplitude_scale = energy_gain / ampl
power_scale     = power / math.pow(ampl,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))