# read data
dat = read_ascii('doc_example2.dat', labels='energy xmu i0')

# do pre-processing steps, here XAFS pre-edge removal
pre_edge(dat.energy, dat.xmu, group=dat)

# select data range to be considered in the fit here.
# note that this could be done inside the objective function,
# but doing these steps here means it done only once.
i1, i2 = index_of(dat.energy, 7105), index_of(dat.energy, 7125)
dat.e = dat.energy[i1+1:i2+1]
dat.y = dat.norm[i1+1:i2+1]

def make_model(pars, data, components=False):
    """make model of spectra: 2 peak functions, 1 erf function, offset"""
    p1 = voigt(data.e, pars.amp1, pars.cen1, pars.wid1)
    p2 = voigt(data.e, pars.amp2, pars.cen2, pars.wid2)
    p3 = voigt(data.e, pars.amp3, pars.cen3, pars.wid3)

    e1 = pars.off + pars.erf_amp * erf( pars.erf_wid*(data.e - pars.erf_cen))
    sum = p1 + p2 + p3 + e1
    if components:
        return sum, p1, p2, p3, e1
    endif
    return sum
enddef

# create a parameter group for the fit:
params = param_group(
    cen1 = param(7113.25, vary=True, min=7111, max=7115),
    cen2 = param(7116.0, vary=True, min=7115, max=7120),

    amp1 = param(0.25, vary=True, min=0),
    amp2 = param(0.25, vary=True, min=0),

    wid1 = param(0.6, vary=True, min=0.05),
    wid2 = param(1.2, vary=True, min=0.05),

    cen3 = param(7122.0, vary=True, min=7120, max=7124),
    amp3 = param(0.5,    vary=True, min=0),
    wid3 = param(1.2,    vary=True, min=0.05),

    off  = param(0.50, vary=True),

    erf_amp = param(0.50, vary=True),
    erf_wid = param(0.50, vary=True),
    erf_cen = param(7123.5, vary=True, min=7121, max=7124)
    )

def resid(pars, data):
    "fit residual"
    return make_model(pars, data) - data.y
enddef

m = minimize(resid, params, args=(dat,))

print( fit_report(m, show_correl=False))

# now plot results (2 different windows)
final, f1, f2, f3, e1 = make_model(params, dat, components=True)
plot(dat.e, dat.y,  label='data', marker='+',
     show_legend=True, legend_loc='ul',
     xlabel='Energy (eV)', ylabel='$\\mu(E)$',
     title='Fe pre-edge peak-fit: best fit and residual',
     new=True)

plot(dat.e, final,  label='fit')
plot(dat.e, (final-dat.y)*10, label='diff (10x)')

plot(dat.e, dat.y, label='data', marker='+',
     show_legend=True, legend_loc='ul',
     xlabel='Energy (eV)', ylabel='$\\mu(E)$',
     title='Fe pre-edge peak-fit: components',
     new=True, win=2)

plot(dat.e, f1, label='peak1', win=2)
plot(dat.e, f2, label='peak2', win=2)
plot(dat.e, e1, label='erf +offset', win=2, color='red', style='dashed')
plot(dat.e, f3, label='peak3', win=2, color='blue', style='dashed')

#
