%matplotlib inline
import numpy as np
import lmfit
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns

import models  # custom module

"""
Explanation: Notebook arguments

sigma (float): standard deviation of additive Gaussian noise to be simulated
time_window (float): seconds, integration window duration
time_step (float): seconds, time step for the moving integration window
time_start (float): seconds, start of time axis (kinetics starts at t = t0).
time_stop (float): seconds, stop of time axis (kinetics starts at t = t0).
t0_vary (bool): whether models should vary the curve origin (t0) during the fit
true_params (dict): parameters used to generate simulated kinetic curves
num_sim_cycles (int): number of times fit is repeated (Monte-Carlo)
taus (tuple): list of values for time-costant tau simulated during repeated fits (Monte-Carlo).


Simulated Kinetic Curve Fit
<p class=lead>This notebook fits simulated exponential transients with additive Gaissian noise in order to study time-constant fitting accuracy.
In particular we compare a simple exponential model with a more realistic model 
with integration window, checking the effect on the fit results.
<p>

You can either run this notebook directly, or run it through the [master notebook](Simulated Kinetic Curve Fit - Run-All.ipynb) for batch processing.

## Imports
End of explanation
"""


labels = ('tau', 'init_value', 'final_value')
model = models.factory_model_exp(t0_vary=True)

"""
Explanation: Fitting models
Models used to fit the data.
1. Simple Exponential
In this model, we define the model function as an exponential transient:
$$ y = f(t) = A \cdot e^{-t/\tau} + K$$
The python function implementing it is:

models.exp_func().

Next cell defines and initializes the fitting model (lmfit.model.Model) including parameters' constrains:
End of explanation
"""


modelw = models.factory_model_expwin(t_window=time_window, decimation=decimation, t0_vary=t0_vary)

"""
Explanation: 2. Integrated Exponential
A more realistic model needs to take into account that each data point 
is the result of an integration over a time window $w$:
$$f(t) = A \cdot e^{-t/\tau} + K$$
$$y(t) = \int_{t}^{t+w} f(t')\;dt'$$
In other words, when we process a measurement in time chunks, we are integrating
a non-stationary signal $f(t)$ over a time window $w$. This integration causes
a smoothing of $f(t)$, regardless of the fact that time is binned or 
is swiped-through with a moving windows (overlapping chunks).
Numerically, $t$ is discretized with step equal to (time_step / decimation).
The python function implementing this model function is:

models.expwindec_func().

And, finally, we define and initialize the fitting model parameters' constrains:
End of explanation
"""


t = np.arange(time_start, time_stop-time_window, time_step).astype(float)
t.size

"""
Explanation: Generative model
These are the models used to generate the simulates (noisy) data.
1. Simple Exponential + Noise
In this simple model, we simulate random data $Y$ as an exponential decay plus
additive Gaussian noise:
$$ Y(t_k) = f(t_k) + N_k $$
$$ {N_k} \sim {\rm Normal}{\mu=0; \sigma}$$
$$ \Delta t = t_k - t_{k-1} = \texttt{time_step}$$
2. Integrated Exponential + Noise
For the "integrating window" model, we first define a finer time axis $\theta_i$ 
which oversamples $t_k$ by a factor $n$. Then we define the function $Y_f$ 
adding Gaussian noise $\sqrt{n}\,N_i$, with $n$ times larger variance:
$$ Y_f(\theta_i) = f(\theta_i) + \sqrt{n}\,N_i $$
$$ \Delta \theta = \theta_i - \theta_{i-1} = \texttt{time_step} \;/\; n$$
Finally, by averaging each time window, we compute the data on the coarse time axis $t_k$:
$$ Y_w(t_k) = \frac{1}{m}\sum_{i} Y_f(\theta_i)$$
Here, for each $t_k$, we compute the mean of $m$ consecutive $Y_f$ values. The number $m$ 
is chosen so that $m\, \Delta \theta$ is equal to the time window.
Noise amplitude
The amplitude of the additive noise ($\sigma$) is estimated from the experimental kinetic curves.
In particular we take the variance from the POST period (i.e. the steady state period after the transient).
The POST period has been chosen because it exhibits higher variance than the PRE period (i.e. the steady state period
before the transient). These values have been calculated in 8-spot bubble-bubble kinetics - Summary.
In both models we define the noise amplitude as sigma (see first cell):
sigma = 0.016

Time axis
We also define the parameters for the time axis $t$:
time_start = -900    # seconds 
time_stop  =  900    # seconds
time_step  = 5       # seconds

Kinetic curve paramenters
The simulated kinetic curve has the following parameters:
true_params = dict(
    tau = 60,          # time constant
    init_value = 0.3,  # initial value (for t &lt; t0) 
    final_value = 0.8, # final value   (for t -&gt; +inf)
    t0 = 0)            # time origin

<div class="alert alert-info">
**NOTE**: All previous parameters are defined in the first notebook cell. 
</div>

Single kinetic curve fit
Here we simulate one kinetic curve and fit it with the two models (simple exponential and integrated exponential).
Draw simulated data
Time axis for simulated data:
End of explanation
"""


y = models.expwindec_func(t, t_window=time_window, **true_params)
y.shape

"""
Explanation: An ideal transient (no noise, no integration):
End of explanation
"""


time_window, time_step

yr = models.expwindec_func(t, t_window=time_window, sigma=sigma, **true_params)
yr.shape

"""
Explanation: A simulated transient (including noise + integration):
End of explanation
"""


plt.plot(t, y, '-', label='model')
plt.plot(t, yr, 'o', label='model + noise')

"""
Explanation: Plot the computed curves:
End of explanation
"""


#%%timeit
resw = modelw.fit(yr, t=t, tau=10, init_value=0.1, final_value=0.9, verbose=False)

"""
Explanation: Fit data
Fit the "Integrated Exponential" model:
End of explanation
"""


#%%timeit
res = model.fit(yr, t=t + 0.5*time_window, tau=10, init_value=0.1, final_value=0.9, verbose=False)

"""
Explanation: Fit the "Simple Exponential" model:
End of explanation
"""


fig = plt.figure(figsize=(14, 8))
res.plot(fig=fig)
ci = lmfit.conf_interval(res, res)
lmfit.report_fit(res)
print(lmfit.ci_report(ci, with_offset=False))
#plt.xlim(-300, 300)

fig = plt.figure(figsize=(14, 8))
resw.plot(fig=fig)
ci = lmfit.conf_interval(resw, resw)
lmfit.report_fit(resw)
print(lmfit.ci_report(ci, with_offset=False))
#plt.xlim(-300, 300)

"""
Explanation: Print and plot fit results:
End of explanation
"""


num_sim_cycles

"""
Explanation: Monte-Carlo Simulation
Here, fixed the model paramenters, we generate and fit several noisy datasets. Then, by plotting the distribution of the fitted parameters, we assess the stability and accuracy of the fit.
Parameters
The number simulation cycles is defined by num_sim_cycles. Current value is:
End of explanation
"""


{k: v for k, v in true_params.items() if k is not "tau"}

"""
Explanation: The fixed kinetic curve parameters are:
End of explanation
"""


taus

t0_vary

"""
Explanation: While tau is varied, taking the following values:
End of explanation
"""


def draw_samples_and_fit(true_params):
    # Create the data
    t = np.arange(time_start, time_stop-time_window, time_step).astype(float)
    yr = models.expwindec_func(t, t_window=time_window, sigma=sigma, decimation=100, **true_params)

    # Fit the model
    tc = t + 0.5*time_window
    kws = dict(fit_kws=dict(nan_policy='omit'), verbose=False)
    res = model.fit(yr, t=tc, tau=90, method='nelder', **kws)
    res = model.fit(yr, t=tc, **kws)
    resw = modelw.fit(yr, t=t, tau=400, decimation=decimation, method='nelder', **kws)
    resw = modelw.fit(yr, t=t, decimation=decimation, **kws)
    return res, resw

def monte_carlo_sim(true_params, N):
    df1 = pd.DataFrame(index=range(N), columns=labels)
    df2 = df1.copy()
    for i in range(N):
        res1, res2 = draw_samples_and_fit(true_params)
        for var in labels:
            df1.loc[i, var] = res1.values[var]
            df2.loc[i, var] = res2.values[var]
    return df1, df2

"""
Explanation: <div class="alert alert-info">
**NOTE**: All previous parameters are defined in the first notebook cell. 
</div>

Functions
Here we define two functions:


draw_samples_and_fit() draws a set of data and fits it with both models


monte_carlo_sim() run the Monte-Carlo simulation: calls draw_samples_and_fit() many times.



NOTE: Global variables are used by previous functions.
End of explanation
"""


mc_results1, mc_results2 = [], []

%%timeit -n1 -r1   # <-- prints execution time
for tau in taus:
    true_params['tau'] = tau
    df1, df2 = monte_carlo_sim(true_params, num_sim_cycles)
    mc_results1.append(df1)
    mc_results2.append(df2)

"""
Explanation: Run Monte-Carlo simulation
Run the Monte-Carlo fit for a set of different time-constants (taus)
and save results in two DataFrames, one for each model.
End of explanation
"""


for tau, df in zip(taus, mc_results1):
    true_params['tau'] = tau
    fig, ax = plt.subplots(1, 3, figsize=(16, 4))
    for i, var in enumerate(labels):
        std = df[var].std()
        df[var].hist(bins=30, ax=ax[i])
        ax[i].set_title("%s = %.1f (%.3f)" % (var, true_params[var], std), fontsize=18)
        ax[i].axvline(true_params[var], color='r', ls='--')
    #print('True parameters: %s' % true_params)

"""
Explanation: <div class="alert alert-danger">
**WARNING**: The previous cell can take a long to execute. Execution time scales with **`num_sim_cycles * len(taus)`**. 
</div>

Results1 - Simple Exponential
End of explanation
"""


for tau, df in zip(taus, mc_results2):
    true_params['tau'] = tau
    fig, ax = plt.subplots(1, 3, figsize=(16, 4))
    for i, var in enumerate(labels):
        std = df[var].std()
        df[var].hist(bins=30, ax=ax[i])
        ax[i].set_title("%s = %.1f (%.3f)" % (var, true_params[var], std), fontsize=18)
        ax[i].axvline(true_params[var], color='r', ls='--')
    #print('True parameters: %s' % true_params)

"""
Explanation: Results2 - Integrated Exponential
End of explanation
"""


