from mathutils import *
from plotutils import *

if is_ui():
    start   = 0.0
    end    = 360.0
    step    = 10.0
    plt = plot(None, title="Harmonic Fit")[0]

#Plot setup
plt.clear()
plt.removeMarker(None)
plt.addSeries(LinePlotSeries("data", color = Color.RED))
plt.setLegendVisible(True)
plt.getAxis(plt.AxisId.X).setRange(start, end)

#Creating an averager of the BPM
avg = create_averager(bpm_x, 3, 0.100)

#The scan command
def after(rec, scan):
    plt.getSeries(0).appendData(rec.positions[0], rec.values[0].mean)
        
r = lscan(phase, [avg], start, end, step, latency=0.2, after_read=after)

#Fitting
rf_phase = r.getPositions(0)
energy   = [val.mean for val in r.getReadable(0)]
(fit_ampl, fit_freq, fit_phase) = fit_harmonic(energy, rf_phase)
ph_crest = round(-fit_phase / fit_freq, 3)
function = HarmonicOscillator(fit_ampl, fit_freq, fit_phase)    
x = frange(min(rf_phase), max(rf_phase), 4.0, True, True)
plot_function(plt, function, "fit", x, show_points=False, color=Color.BLUE)
plt.addMarker(ph_crest, None, "ph_crest="+str(ph_crest), Color.MAGENTA)

#Return values
print "ph_crest = ", ph_crest
phase.write(ph_crest)
phase_offset    =  round(90.0 - ph_crest, 3)
amplitude_scale =  round(fit_ampl / amplitude.read(), 3)
power_scale     =  round(power.read() / math.pow(fit_ampl,2), 3)
set_return([phase_offset, amplitude_scale,power_scale ])