from saba import SherpaFitter
from astropy.modeling.models import Gaussian1D, Gaussian2D

import numpy as np
import matplotlib.pyplot as plt

sfitter = SherpaFitter(statistic='chi2', optimizer='levmar', estmethod='confidence')
sfitter.est_config['max_rstat'] = 4
np.random.seed(0x1337)

true = Gaussian1D(amplitude=3, mean=0.9, stddev=0.5)
err = 0.05
step = 0.1
x = np.arange(-3, 3, step)
y = true(x) + err * np.random.uniform(-1, 1, size=len(x))

yerrs = err * np.random.uniform(0.2, 1, size=len(x))
binwidth = step * np.ones(x.shape)

fit_model = true.copy()  # ofset fit model from true
fit_model.amplitude = 2
fit_model.mean = 0
fit_model.stddev = 0.2

fitted_model = sfitter(fit_model, x, y, xbinsize=binwidth, err=yerrs)

param_errors = sfitter.est_errors(sigma=3)
min_model = fitted_model.copy()
max_model = fitted_model.copy()

for pname, pval, pmin, pmax in list(zip(*param_errors)):
    getattr(min_model, pname).value = pval + pmin
    getattr(max_model, pname).value = pval + pmax

plt.plot(x, true(x), label="True")
plt.errorbar(x, y, xerr=binwidth, yerr=yerrs, ls="")
plt.plot(x, fitted_model(x), label="Fitted model")
plt.plot(x, min_model(x), label="min model", ls="--")
plt.plot(x, max_model(x), label="max model", ls="--")
plt.legend(loc=2, frameon=False)
plt.xlim((-3, 3))
plt.show()