import numpy as np

from astropy.models import models, fitting, parameters

# We use iminuit to check the results
# http://iminuit.github.io/iminuit/
# iminuit is super light-weight, all you get is a class called Minuit,
# which can compute best-fit parameters (migrad), parameter errors
# (hesse and minos) and a few other things (e.g. fit statistic 
# profiles / contours).
# You have to write your models and fit statistics yourself.
from iminuit import Minuit

%pylab inline

# Data
# Let's use the example from here:
# http://nbviewer.ipython.org/5030045
xdata = np.array([0, 1, 2, 3, 4, 5])
ydata = np.array([1, 1, 5, 7, 8, 12])
sigma = np.array([1, 2, 1, 2, 1, 2], dtype=float)

# Plot the data
plt.plot(xdata, ydata, 'o', label='true model')
plt.xlim(-0.5, 5.5)
plt.ylim(0, 13);

# Model
def line(x, slope, intercept):
    """Linear function (two fit parameters)"""
    return slope * x + intercept

# Fit statistic (not used at the moment)
def chi2(model, parameters, data):
    """Generic chi^2 function"""
    slope, intercept = parameters
    x, y, sigma = data
    chi = (y - model(*parameters)) / sigma
    chi2 = chi ** 2
    return np.sum(chi2)

def chi2_for_line(slope, intercept):
    """
    Chi^2 statistic for given line model parameters.
    Note: The Minuit fitter is not interested in your data or model,
    only this function!
    """
    chi = (ydata - line(xdata, slope, intercept)) / sigma
    chi2 = chi ** 2
    return np.sum(chi2)

minuit = Minuit(chi2_for_line, pedantic=False, print_level=0)
minuit.migrad() # Compute optimum values
minuit.hesse() # Compute errors (covariance matrix as inverse Hesse matrix)
print 'values:', minuit.values
print 'errors:', minuit.errors

# Compute parameters and errors with astropy.models
poly_model = models.Poly1DModel(1)
fitter = fitting.LinearLSQFitter(poly_model)
fitter(xdata, ydata, weights=sigma ** (-1))
print dict(slope=poly_model.c1, intercept=poly_model.c0)

# TODO: Currently no way to compute parameter errors with astropy.models.

class UserLineModel(models.PModel):
    """
    Line as a user-defined model in astropy

    Note: this is a lot of boilerplate code.
    But in reality there is no reason for a user
    to not use the built-in polynomial models,
    so it doesn't matter if this is not easy in astropy.
    """
    parnames = ['slope', 'intercept']

    def __init__(self, slope, intercept, paramdim=1):
        self.linear = True 
        self._slope = parameters._Parameter(name='slope', val=slope, mclass=self, paramdim=paramdim)
        self._intercept = parameters._Parameter(name='intercept', val=intercept, mclass=self, paramdim=paramdim)
        models.ParametricModel.__init__(self, self.parnames, ndim=1, outdim=1, paramdim=paramdim)
        self.domain = None
        self.window = None
        self._order = 2

    def eval(self, x, params):
        return params[0] * x + params[1]

    def __call__(self, x):
        x, format = models._convert_input(x, self.paramdim)
        result = self.eval(x, self.psets)
        return models._convert_output(result, format)
    
    def deriv(self, x):
        res = np.ones((len(x), 2))
        res[:,0] = x.copy()
        return res

# Run the example using the user-defined model
line_model = UserLineModel(slope=1, intercept=1)
fitter = fitting.LinearLSQFitter(line_model)
fitter(xdata, ydata, weights=sigma ** (-1))
print dict(slope=line_model.slope, intercept=line_model.intercept)
# TODO: Results don't match. Why!???

# Trying to find out why results don't match ...
poly_model.__dict__
line_model.__dict__

MAXITER = 100
EPS = 1e-10

class UserSLSQPFitter(fitting.Fitter):
    """
    User-defined chi^2 fitter in astropy

    Note: this is a lot of boilerplate code.
    """
    def __init__(self, model, fixed=None, tied=None, bounds=None,
                 eqcons=None, ineqcons=None):
        fitting.Fitter.__init__(self, model, fixed=fixed, tied=tied, bounds=bounds,
                                eqcons=eqcons, ineqcons=ineqcons)

    def errorfunc(self, fps, *args):
        meas = args[0]
        self.fitpars = fps
        res = self.model(*args[1:]) - meas
        return np.sum(res**2)

    def __call__(self, x, y , maxiter=MAXITER, epsilon=EPS):
        self.fitpars = optimize.fmin_slsqp(self.errorfunc, p0=self.model.parameters[:], args=(y, x),
        bounds=self.constraints._bounds, eqcons=self.constraints.eqcons,
        ieqcons=self.constraints.ineqcons)


# Run an example with the user-defined chi^2 fitter
# TODO: Make this work
model = models.Poly1DModel(1)
fitter = UserSLSQPFitter(model)

# TODO: Run an example with the user-defined model and chi^2 fitter!

# Model
def gauss(x, amplitude, mean, sigma):
    exponent = 0.5 * ((x - mean) / sigma) ** 2
    return amplitude * np.exp(-exponent)

# Data (simulated)
true_parameter_values=dict(amplitude=100, mean=5, sigma=1)
xdata = np.arange(0, 10, 0.1)
ydata_model = gauss(xdata, **true_parameter_values)
sigma = 10 * np.ones_like(xdata)
np.random.seed(0)
yerror = np.random.normal(0, sigma)
ydata = ydata_model + yerror

# Plot the data
plt.plot(xdata, ydata_model, lw=2, label='true model')
plt.plot(xdata, ydata, 'o', label='simulated data')
plt.legend();

def chi2_for_gauss(amplitude, mean, sigma):
    """
    Chi^2 statistic for given gauss model parameters.
    """
    chi = (ydata - gauss(xdata, amplitude, mean, sigma)) / sigma
    chi2 = chi ** 2
    return np.sum(chi2)

minuit = Minuit(chi2_for_gauss, pedantic=False, print_level=0, **true_parameter_values)
minuit.migrad() # Compute optimum values
minuit.hesse() # Compute errors (covariance matrix as inverse Hesse matrix)
print 'values:', minuit.values
print 'errors:', minuit.errors

# Compute parameters and errors with astropy.models
gauss_model = models.Gauss1DModel(amplitude=true_parameter_values['amplitude'],
                                  xcen=true_parameter_values['mean'],
                                  xsigma=true_parameter_values['sigma'])
fitter = fitting.NonLinearLSQFitter(gauss_model)
fitter(xdata, ydata, weights=sigma ** (-1))
print dict(amplitude=gauss_model.amplitude,
           mean=gauss_model.xcen,
           sigma=gauss_model.xsigma)

# TODO: Results don't match the iminuit results from the last cell. Why?

# TODO: User-defined Gauss1D model

# TODO: User-defined NonLinearLSQFitter

# Model
def gauss(x, amplitude, mean, sigma):
    exponent = 0.5 * ((x - mean) / sigma) ** 2
    return amplitude * np.exp(-exponent)

# Data (simulated)
true_parameter_values=dict(amplitude=100, mean=5, sigma=1)
xdata = np.arange(0, 10, 0.1)
ydata_model = gauss(x, **true_parameter_values)
np.random.seed(0)
ydata = np.random.poisson(ydata_model)

# Plot the data
plt.plot(xdata, ydata_model, lw=2, label='true model')
plt.plot(xdata, ydata, 'o', label='simulated data')
plt.legend();

# Fit statistic
def cash(D, M, safe=False):
    """cash fit statistic where D = data, M = model
    http://cxc.cfa.harvard.edu/sherpa/statistics/#cash

    This is simply 2 x the negative log Poisson likelihood:
    P = (M ** D) * exp(-M) / (D!)
    log(P) = D * log(M) - M - log(D!)
    cash = - 2 * log(P) (model-independent term dropped)
    """
    stat = 2 * (M - D * log(M))
    if safe:
        return np.where(M > 0, stat, 1)
    else:
        return stat

def cash_for_gauss(amplitude, mean, sigma):
    """
    Cash statistic for given gauss model parameters.
    Note: The Minuit fitter is not interested in your data or model,
    only this function!
    """
    ydata_model = gauss(xdata, amplitude, mean, sigma)
    statval = cash(ydata, ydata_model)
    return np.sum(statval)

# TODO: check if an up=0.5 is missing to get correct errors!
minuit = Minuit(cash_for_gauss, pedantic=False, print_level=0, **true_parameter_values)
minuit.migrad() # Compute optimum values
minuit.hesse() # Compute errors (covariance matrix as inverse Hesse matrix)
print 'values:', minuit.values
print 'errors:', minuit.errors

# TODO: reproduce Minuit results from last cell with astropy

# TODO: User-defined cash fitter
'''
class UserCashFitter(fitting.Fitter):
    """User-defined cash statistic fitter in astropy"""
    def __init__(self, model, fixed=None, tied=None, bounds=None,
                 eqcons=None, ineqcons=None):
        fitting.Fitter.__init__(self, model, fixed=fixed, tied=tied, bounds=bounds,
                                eqcons=eqcons, ineqcons=ineqcons)

    def errorfunc(self, fps, *args):
        meas = args[0]
        self.fitpars = fps
        res = self.model(*args[1:]) - meas
        return np.sum(res**2)

    def __call__(self, x, y , maxiter=MAXITER, epsilon=EPS):
        self.fitpars = optimize.fmin_slsqp(self.errorfunc, p0=self.model.parameters[:], args=(y, x),
        bounds=self.constraints._bounds, eqcons=self.constraints.eqcons,
        ieqcons=self.constraints.ineqcons)
'''

'''
class OldLine(models.ParametricModel):
    """Old version of the Line class. Doesn't work."""
    parnames = ['slope', 'intercept']
    def __init__(self, slope, intercept, paramdim=1):
        self.linear = True
        self._slope = parameters._Parameter(name='slope', val=slope, mclass=self, paramdim=paramdim)
        self._intercept = parameters._Parameter(name='intercept', val=intercept, mclass=self, paramdim=paramdim)
        models.ParametricModel.__init__(self, self.parnames, ndim=1, outdim=1, paramdim=paramdim)
    def eval(self, x, params):
        return params[0] * x + params[1]
    def __call__(self, x):
        x, format = models._convert_input(x, self.paramdim)
        result = self.eval(x, self.psets)
        return models._convert_output(result, format)
'''