# AUTOGENERATED! DO NOT EDIT! File to edit: 03_fit.ipynb (unless otherwise specified).

__all__ = ['Parameter', 'residual', 'fit']

# Cell
import lmfit

# Cell
import numpy as np
from lmfit import minimize, Parameters, report_fit


# Cell
def Parameter(name,**kwargs):
    from lmfit import Parameters
    params = Parameters()
    params.add(name, **kwargs)

    return params


# Cell
def residual(ps, sim):

    params={}
    for key in ps.keys():
        if key.startswith('initial_'):
            name=key.split('initial_')[1]
            _c=sim.get_component(name)
            _c.initial_value=ps[key].value

        else:
            params[key]=ps[key].value

    sim.params(**params)

    # run the sim
    sim.run_fast()

    # compare with data
    values=[]
    for _c in sim.components+sim.assignments:
        if not _c.data:
            continue
        t=np.array(_c.data['t']).ravel()
        y=np.array(_c.data['value']).ravel()
        y_fit=sim.interpolate(t,_c.name)

        if any(np.isnan(y_fit)):
            values.append([-np.inf])
        elif any(abs(y_fit)>1e100):
            values.append([-np.inf])
        else:
            values.append(y-y_fit)

    return np.concatenate(values).ravel()


def fit(sim,
       *args,method='leastsq'):

    found_data=False
    for _c in sim.components+sim.assignments:
        if not _c.data:
            continue
        found_data=True

    if not found_data:
        raise ValueError("No data given for the simulation...can't fit.")



    from lmfit import Parameters,minimize

    fitparams=Parameters()
    for arg in args:
        fitparams+=arg

    for key in fitparams.keys():
        if key.startswith('initial_'):
            name=key.split('initial_')[1]
            try:
                _c=sim.get_component(name)
            except IndexError:
                raise ValueError("%s is a bad initial variable because %s is not a variable in the dynamical model." % (key,name))
        else:
            if not key in sim.original_params:
                raise ValueError("%s is not a parameter in the dynamical model.  Parameters are %s" % (key,str(sim.original_params)))


    result = minimize(residual, fitparams, args=(sim,), method=method)

    params={}
    ps=result.params
    for key in ps.keys():
        if key.startswith('initial_'):
            name=key.split('initial_')[1]
            _c=sim.get_component(name)
            _c.initial_value=ps[key].value

        else:
            params[key]=ps[key].value

    sim.params(**params)


    return result

