doc =  """
   Example fit to Gaussian line shape.
"""

show_plots = True
import time
t0 = time.time()

xlog2 = log(2.0)
s2pi = sqrt(xlog2/pi)
parformat = " %4i  % .4f  % .4f  % .4f  % .4f % .4f  % .4f % .4f"

def mymodel(x, _off, _amp, _cen, _wid):
    """standard gaussian with scale and constant offset """
    return _off + _amp*(s2pi / _wid) * exp(-xlog2 * (x-_cen)**2 / _wid**2)
#enddef

def resid(pars, data, verbose=False, doplot=False, update_rate=5):
    model = mymodel(data.x, pars.off1, pars.amp1, pars.cen1, pars.wid1) + \
            mymodel(data.x, pars.off2, pars.amp2, pars.cen2, pars.wid2)
    if verbose:
        if pars.fit_iter == 0:
            print( "    I   Offset    Amp1      Cen1    Sigma1   Amp2     Cen2    Sigma2")
        #endif
        if pars.fit_iter % update_rate == 0:
            print( parformat % (pars.fit_iter, pars.off1.value,
                               pars.amp1.value, pars.cen1.value,
                               pars.wid1.value, pars.amp2.value,
                               pars.cen2.value, pars.wid2.value))
	    if doplot: update_trace(data.x, model, trace=2)
        #endif
    #endif
    pars.fit_iter += 1
    return  model - data.var
#enddef

# create data
xdat = group(x=arange(501)/20.0)
xdat.var =  mymodel(xdat.x, -0.2, 15.0, 11.0, .5) + \
            mymodel(xdat.x,  0.0, 12.0, 9.0, 2.0) + \
            random.normal(size=len(xdat.x), scale=.20)

# create group for parameters
fitpars = param_group(off1 = param(0, vary=True),
                cen1 = param(10.0, vary=True),
                amp1 = param(30.0, vary=True),
                wid1 = param(2.0, vary=True),

                off2 = param(0.0, vary=False),
                cen2 = param(expr='cen1 - 2'),
                amp2 = param(50.0, vary = True),
                wid2 = param(expr='wid1*4'),
                fit_iter=0)

if show_plots:
    init = mymodel(xdat.x, fitpars.off1, fitpars.amp1, fitpars.cen1, fitpars.wid1) + \
           mymodel(xdat.x, fitpars.off2, fitpars.amp2, fitpars.cen2, fitpars.wid2)
    newplot(xdat.x, xdat.var, color='blue')
    plot(xdat.x, init, color='red')
    plot(xdat.x, init, color='black', style='--')
endif

minout = minimize(resid, fitpars, args=(xdat,),
                  kws={'verbose':show_plots,
                       'update_rate': 10,
                       'doplot': show_plots})

if show_plots:
    final = mymodel(xdat.x, fitpars.off1, fitpars.amp1, fitpars.cen1, fitpars.wid1) + \
            mymodel(xdat.x, fitpars.off2, fitpars.amp2, fitpars.cen2, fitpars.wid2)
    update_trace(xdat.x, final, trace=2)
endif
print( fit_report(minout))
print( 'Total time for fit = %.3f sec' % (time.time() - t0))

# print( 'Calculating Chi-Square map for amp2 and off1:')
#
# # important!!! supress plots, reset counter
# minout.userkws['plot'] = False
# minout.userkws['update_rate'] = 1000
# fitpars.fit_iter = 0
#
# conf_int = confidence_intervals(minout)
# print( confidence_report(conf_int))
#
# xvals, yvals, c2map = chi2_map(minout, 'amp1', 'off1', nx=21, ny=11, xrange=[13,17], yrange=[-0.4, -0.0])
# contour(c2map, x=xvals, y=yvals, title='Chi2 map for Amp1 and Off1')
