#!/usr/bin/env python
# coding: utf-8

# # Live Nonlinear Fitting
# 
# ## Configuration
# 
# This code would normally go in a script automatically run at startup. The user would not have to worry about this.

# In[ ]:


get_ipython().run_line_magic('matplotlib', 'widget')
import matplotlib.pyplot as plt

import numpy as np
import lmfit
from bluesky import RunEngine
from bluesky.plans import scan
from ophyd.sim import motor, noisy_det
from bluesky.callbacks import LiveFit, LivePlot, LiveFitPlot
import matplotlib.pyplot as plt

RE = RunEngine({})


# ## Data Acquisition

# In[ ]:


fig, ax = plt.subplots() # As before, create plot to be able to follow the progress of the fit during the scan


# In[ ]:


# define Gaussian function and save it as a model using lmfit
def gaussian(x, A, sigma, x0):
    return A*np.exp(-(x - x0)**2/(2 * sigma**2))

model = lmfit.Model(gaussian)

# The initial guess (gray on the plot) is just a starting point from which the fitting algorithm works.
init_guess = {'A': 2,
              'sigma': lmfit.Parameter('sigma', 3, min=0),
              'x0': -0.2}


# In[ ]:


# LiveFit example
lf = LiveFit(model, 'noisy_det', {'x': 'motor'}, init_guess)

# now add the plot...
lfp = LiveFitPlot(lf, color='r', ax=ax)

RE(scan([noisy_det], motor, -1, 1, 100), lfp)


# In[ ]:


plt.gcf() # Unfortunately, since we do not see the motor positions relative to the fit,
          # we cannot evaluate how well the fit worked.


# In[ ]:


# Show motor position vs. fit during the scan

fig, ax = plt.subplots()  # explitly create figure to follow fitting


# In[ ]:


lfp = LiveFitPlot(lf, ax=ax, color='r')
lp = LivePlot('noisy_det', 'motor', ax=ax, marker='o', linestyle='none') # plotting the noisy_det scalar object

RE(scan([noisy_det], motor, -1, 1, 100), [lfp, lp])


# In[ ]:


# We can see that from imperfect initial conditions, the fit now matches the data well
plt.gcf() # Display a snapshot of the current state of the figure.


# In[ ]:


# Get the final fit parameters
lf.result.best_values


# ## Exercises
# 
# 1. Try a different user-defined peak-like function to use as a model.