import ch.psi.pshell.epics.ChannelDouble as ChannelDouble
from mathutils import fit_polynomial, PolynomialFunction
A1 = ChannelDouble("K value",                  "SARUN08-UIND030:K_SET")
#S1 = ChannelDouble("Energy per pulse (uJ)",    "SARFE10-PBPG050:PHOTON-ENERGY-PER-PULSE-US")
#S1 = ChannelDouble("Energy per pulse (uJ)",    "SARFE10-PBPG050:HAMP-INTENSITY-CAL")
S1 = ChannelDouble("Hamp RAW",    "SARFE10-PBIG050-EVR0:CALCI")
A1.initialize()
S1.initialize()
A1_init = A1.read()
A1i    = A1_init - 0.005
A1f    = A1_init + 0.005
nstep  = 10
lat    = 0.01
nav    = 100
wait   = 3
plt    = plot(None, title="Output")[0]
plt.clear()
plt.setStyle(plt.Style.ErrorY)
plt.addSeries(LinePlotErrorSeries("Sensor1", Color.red))

def after_sample(record, scan):
    plt.getSeries(0).appendData(record.positions[0], record.readables[0].mean, record.readables[0].stdev)
    
try:
    S1_averager = create_averager(S1, nav, lat)
    A1.write(A1i)
    time.sleep(wait)
    r = lscan(A1, (S1_averager), A1i, A1f, nstep, latency=5.0, after_read = after_sample)
    Act1     = r.getPositions(0)
    S1mean   = [val.mean for val in r.getReadable(0)]
    S1rms    = [val.stdev for val in r.getReadable(0)]
finally:
    A1.write(A1_init)
    A1.close()
    S1.close()

    ## add fitting:
    pars_polynomial = fit_polynomial(S1mean, Act1, 3)
    p1 = PolynomialFunction(pars_polynomial)
    resolution = (A1f - A1i)/100
    fit_polinomial = []
    for x in frange(A1i, A1f, resolution, True):
        fit_polinomial.append(p1.value(x))
    x = frange(A1i, A1f+resolution, resolution)
    #plot(x, fit_polinomial)
    plots = plot([S1mean, fit_polinomial] , 
    ["data", "polinomial"], xdata = [Act1,x], title="Data")