# AUTOGENERATED! DO NOT EDIT! File to edit: ../03_fit.ipynb.

# %% auto 0
__all__ = ['Parameter', 'residual', 'fit']

# %% ../03_fit.ipynb 3
import lmfit

# %% ../03_fit.ipynb 6
import numpy as np
from lmfit import minimize, Parameters, report_fit

# %% ../03_fit.ipynb 11
def Parameter(name,**kwargs):
    from lmfit import Parameters
    params = Parameters()
    params.add(name, **kwargs)
    
    return params

# %% ../03_fit.ipynb 12
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

