measurement_id = 0
windows = (60, 180)

# Cell inserted during automated execution.
windows = (30, 180)
measurement_id = 0

"""
Explanation: Executed: Tue Mar 28 00:43:07 2017
Duration: 33 seconds.
End of explanation
"""


import time
from pathlib import Path
import pandas as pd
from scipy.stats import linregress
from scipy import optimize
from IPython.display import display

from fretbursts import *

sns = init_notebook(fs=14)

import lmfit; lmfit.__version__

import phconvert; phconvert.__version__

"""
Explanation: Notebook arguments

measurement_id (int): Select the measurement from the list. Valid values: 0 .. 3

1-spot realtime kinetics
<p class=lead>This notebook executes the realtime-kinetics analysis.</p>

The first cell of this notebook selects which measurement is analyzed. 
Measurements can be processed one-by-one, by manually running this notebook,
or in batch by using the notebook: "1-spot bubble-bubble kinetics - Run-All".
Loading the software
End of explanation
"""


path = Path('./data/')
pattern = 'singlespot*.hdf5'
filenames = list(str(f) for f in path.glob(pattern))
filenames

basenames = list(f.stem for f in path.glob(pattern))
basenames

start_times = [600, 900, 900,
               600, 600, 600, 600, 600, 600, 
               600, 600, 600]  # time of NTP injection and start of kinetics

filename = filenames[measurement_id]
start_time = start_times[measurement_id]

filename

import os
assert os.path.exists(filename)

"""
Explanation: Selecting a data file
End of explanation
"""


d = loader.photon_hdf5(filename)

plot_alternation_hist(d)

loader.alex_apply_period(d)

d.time_max

"""
Explanation: Data load and Burst search
Load and process the data:
End of explanation
"""


d.calc_bg(bg.exp_fit, time_s=10, tail_min_us='auto', F_bg=1.7)

"""
Explanation: Compute background and burst search:
End of explanation
"""


dplot(d, hist_bg);

"""
Explanation: Let's take a look at the photon waiting times histograms and at the fitted background rates:
End of explanation
"""


dplot(d, timetrace_bg);
xlim(start_time - 150, start_time + 150)

"""
Explanation: Using dplot exactly in the same way as for the single-spot data has now generated 8 subplots, one for each channel.
Let's plot a timetrace for the background to see is there are significat variations during the measurement:
End of explanation
"""


#dplot(d, timetrace)
#xlim(2, 3); ylim(-100, 100);

"""
Explanation: We can look at the timetrace of the photon stream (binning):
End of explanation
"""


#%%timeit -n1 -r1
ddc = bext.burst_search_and_gate(d)

ds1 = ddc.select_bursts(select_bursts.size, th1=25)
ds = ds1.select_bursts(select_bursts.naa, th1=25)

"""
Explanation: Burst selection and FRET
End of explanation
"""


bpl.alex_jointplot(ds)

ds0 = ds.select_bursts(select_bursts.time, time_s1=0, time_s2=start_time-10)

dplot(ds0, hist_fret, pdf=False);

weights = 'size'
bext.bursts_fitter(ds0, weights=weights)
ds0.E_fitter.fit_histogram(mfit.factory_two_gaussians(p1_center=0.5, p2_center=0.9), verbose=False)
dplot(ds0, hist_fret, show_model=True, weights=weights);
ds0.E_fitter.params

weights = None
bext.bursts_fitter(ds0, weights=weights)
ds0.E_fitter.fit_histogram(mfit.factory_two_gaussians(p1_center=0.5, p2_center=0.9), verbose=False)
dplot(ds0, hist_fret, show_model=True, weights=weights);
ds0.E_fitter.params

"""
Explanation: Selecting bursts by size
End of explanation
"""


def gauss2(**params0):
    peak1 = lmfit.models.GaussianModel(prefix='p1_')
    peak2 = lmfit.models.GaussianModel(prefix='p2_')
    model = peak1 + peak2
    model.set_param_hint('p1_center', **{'value': 0.6, 'min': 0.3, 'max': 0.8, **params0.get('p1_center', {})})
    model.set_param_hint('p2_center', **{'value': 0.9, 'min': 0.8, 'max': 1.0, **params0.get('p2_center', {})})
    for sigma in ['p%d_sigma' % i for i in (1, 2)]:
        model.set_param_hint(sigma, **{'value': 0.02, 'min': 0.01, **params0.get(sigma, {})})
    for ampl in ['p%d_amplitude' % i for i in (1, 2)]:
        model.set_param_hint(ampl, **{'value': 0.5, 'min': 0.01, **params0.get(ampl, {})})
    model.name = '3 gauss peaks'
    return model

#%matplotlib notebook

#fig, ax = plt.subplots(figsize=(12, 8))
#dplot(dm0, scatter_fret_size, ax=ax)

bext.bursts_fitter(ds0, weights=None)
ds0.E_fitter.fit_histogram(gauss2(), verbose=False)
mfit.plot_mfit(ds0.E_fitter)
params_2gauss = ds0.E_fitter.params
plt.xlabel('E')
plt.ylabel('PDF')
plt.title('')
params_2gauss

ds_final = ds.select_bursts(select_bursts.time, time_s1=start_time+300, time_s2=ds.time_max + 1)
ds_final.num_bursts

bext.bursts_fitter(ds_final, weights=None)
model = gauss2()
model.set_param_hint('p2_center', value=params_2gauss.p2_center[0], vary=False)
ds_final.E_fitter.fit_histogram(model, verbose=False)
fig, ax = plt.subplots(figsize=(12, 6))
mfit.plot_mfit(ds_final.E_fitter, ax=ax)
params_2gauss1 = ds_final.E_fitter.params
params_2gauss1

#del params_2gauss0

is_runoff = 'runoff' in filename.lower()

if 'params_2gauss0' not in locals():
    params_2gauss0 = params_2gauss.copy()
    if is_runoff:
        params_2gauss0.p2_center = params_2gauss1.p2_center
    else:
        params_2gauss0.p1_center = params_2gauss1.p1_center

params_2gauss0.p1_amplitude + params_2gauss0.p2_amplitude

'params_2gauss0' in locals()

"""
Explanation: 2-Gaussian peaks
End of explanation
"""


from scipy import optimize

params_fixed = dict(
    mu1=float(params_2gauss0.p1_center),
    mu2=float(params_2gauss0.p2_center),
    sig1=float(params_2gauss0.p1_sigma),
    sig2=float(params_2gauss0.p2_sigma),
)

def em_weights_2gauss(x, a2, mu1, mu2, sig1, sig2):
    """Responsibility function for a 2-Gaussian model.
    
    Return 2 arrays of size = x.size: the responsibility of 
    each Gaussian population.
    """
    a1 = 1 - a2
    assert np.abs(a1 + a2 - 1) < 1e-3
    f1 = a1 * gauss_pdf(x, mu1, sig1)
    f2 = a2 * gauss_pdf(x, mu2, sig2)
    γ1 = f1 / (f1 + f2)
    γ2 = f2 / (f1 + f2)
    return γ1, γ2

def em_fit_2gauss(x, a2_0, params_fixed, print_every=10, max_iter=100, rtol=1e-3):
    a2_new = a2_0
    rel_change = 1
    i = 0
    while rel_change > rtol and i < max_iter:

        # E-step
        γ1, γ2 = em_weights_2gauss(x, a2_new, **params_fixed)
        assert np.allclose(γ1.sum() + γ2.sum(), x.size)

        # M-step
        a2_old = a2_new
        a2_new = γ2.sum()/γ2.size

        # Convergence
        rel_change = np.abs((a2_old - a2_new)/a2_new)
        i += 1
        if (i % print_every) == 0:
            print(i, a2_new, rel_change)
        
    return a2_new, i

from matplotlib.pylab import normpdf as gauss_pdf

# Model PDF to be maximized
def model_pdf(x, a2, mu1, mu2, sig1, sig2):
    a1 = 1 - a2
    #assert np.abs(a1 + a2 + a3 - 1) < 1e-3
    return (a1 * gauss_pdf(x, mu1, sig1) + 
            a2 * gauss_pdf(x, mu2, sig2))

def func2min_lmfit(params, x):
    a2 = params['a2'].value
    mu1 = params['mu1'].value
    mu2 = params['mu2'].value
    sig1 = params['sig1'].value
    sig2 = params['sig2'].value
    return -np.sqrt(np.log(model_pdf(x, a2, mu1, mu2, sig1, sig2)))

def func2min_scipy(params_fit, params_fixed, x):
    a2 = params_fit
    mu1 = params_fixed['mu1']
    mu2 = params_fixed['mu2']
    sig1 = params_fixed['sig1']
    sig2 = params_fixed['sig2']
    return -np.log(model_pdf(x, a2, mu1, mu2, sig1, sig2)).sum()

# create a set of Parameters
params = lmfit.Parameters()
params.add('a2', value=0.5, min=0)
for k, v in params_fixed.items():
    params.add(k, value=v, vary=False)

"""
Explanation: Fit
End of explanation
"""


x = ds0.E_
#x

#result = lmfit.minimize(func2min_lmfit, params, args=(x,), method='nelder')
#lmfit.report_fit(result.params)

#optimize.brute(func2min_scipy, ranges=((0.01, 0.99), (0.01, 0.99)), Ns=101, args=(params, x))

res_em = em_fit_2gauss(x, 0.5, params_fixed)
res_em

res = optimize.minimize(func2min_scipy, x0=[0.5], args=(params_fixed, x), method='Nelder-Mead')
res

res = optimize.minimize(func2min_scipy, x0=[0.5], args=(params_fixed, x), bounds=((0,1),), method='SLSQP')
res

res = optimize.minimize(func2min_scipy, x0=[0.5], args=(params_fixed, x), bounds=((0,1),), method='TNC')
res

bins = np.arange(-0.1, 1.1, 0.025)
plt.hist(x, bins, histtype='step', lw=2, normed=True);
xx = np.arange(-0.1, 1.1, 0.005)
#plt.plot(xx, model_pdf(xx, params))
plt.plot(xx, model_pdf(xx, a2=res_em[0], **params_fixed))

"""
Explanation: $$f(x) = \frac{A}{\sigma\sqrt{2\pi}}\, e^{-\frac{(x - \mu)^2}{2 \sigma^2}}$$
$$\log f(x) = \log \frac{A}{\sigma\sqrt{2\pi}}\, e^{-\frac{(x - \mu)^2}{2 \sigma^2}} = \log{A} -\log{\sigma} - \log\sqrt{2\pi} -\frac{(x - \mu)^2}{2 \sigma^2}$$
$$w_1 \; f_1(x) + w_2 \; f_2(x) + w_3 \; f_3(x)$$
$$\log (w_1 \; f_1(x)) = \log{w_1} + \log{f_1(x)}$$
End of explanation
"""


def _kinetics_fit_em(dx, a2_0, params_fixed, **kwargs):
    kwargs = {'max_iter': 100, 'print_every': 101, **kwargs}
    a2, i = em_fit_2gauss(dx.E_, a2_0, params_fixed, **kwargs)
    return a2, i < kwargs['max_iter']

def _kinetics_fit_ll(dx, a2_0, params_fixed, **kwargs):
    kwargs = {'method':'Nelder-Mead', **kwargs}
    res = optimize.minimize(func2min_scipy, x0=[a2_0], args=(params_fixed, dx.E_), 
                            **kwargs)
    return res.x[0], res.success
    
def _kinetics_fit_hist(dx, a2_0, params_fixed):
    E_fitter = bext.bursts_fitter(dx)
    model = mfit.factory_two_gaussians()
    model.set_param_hint('p1_center', value=params_fixed['mu1'], vary=False)
    model.set_param_hint('p2_center', value=params_fixed['mu2'], vary=False)
    model.set_param_hint('p1_sigma', value=params_fixed['sig1'], vary=False)
    model.set_param_hint('p2_sigma', value=params_fixed['sig2'], vary=False)
    E_fitter.fit_histogram(model, verbose=False)
    return (float(E_fitter.params.p2_amplitude), 
            dx.E_fitter.fit_res[0].success)

def kinetics_fit(ds_slices, params_fixed, a2_0=0.5, method='em', **method_kws):
    fit_func = {
        'em': _kinetics_fit_em, 
        'll': _kinetics_fit_ll,
        'hist': _kinetics_fit_hist}
    
    fit_list = []
    for dx in ds_slices:
        a2, success = fit_func[method](dx, a2_0, params_fixed, **method_kws)
        df_i = pd.DataFrame(data=dict(p2_amplitude=a2,
                                      p1_center=params_fixed['mu1'], p2_center=params_fixed['mu2'], 
                                      p1_sigma=params_fixed['sig1'], p2_sigma=params_fixed['sig2'],
                                      tstart=dx.slice_tstart, tstop=dx.slice_tstop, 
                                      tmean=0.5*(dx.slice_tstart + dx.slice_tstop)), 
                            index=[0.5*(dx.slice_tstart + dx.slice_tstop)])
        if not success:
            print('* ', end='', flush=True)
            continue
        
        fit_list.append(df_i)
    print(flush=True)
    return pd.concat(fit_list)

start_time/60

"""
Explanation: Kinetics
Definitions
End of explanation
"""


def print_slices(moving_window_params):
    msg = ' - Slicing measurement:'
    for name in ('start', 'stop', 'step', 'window'):
        msg += ' %s = %.1fs' % (name, moving_window_params[name]) 
    print(msg, flush=True)
    num_slices = len(bext.moving_window_startstop(**moving_window_params))
    print('   Number of slices %d' % num_slices, flush=True)

t1 = time.time()
time.ctime()

ds.calc_max_rate(m=10)
ds_high = ds.select_bursts(select_bursts.E, E1=0.85)

step = 10
params = {}
for window in windows:
    moving_window_params = dict(start=0, stop=ds.time_max, step=step, window=window)
    print_slices(moving_window_params)

    ds_slices = bext.moving_window_chunks(ds, time_zero=start_time, **moving_window_params)
    for meth in ['em', 'll', 'hist']:
        print('    >>> Fitting method %s ' % meth, end='', flush=True)
        p = kinetics_fit(ds_slices, params_fixed, method=meth)
        print(flush=True)
        p['kinetics'] = p.p2_amplitude
        p = p.round(dict(p1_center=3, p1_sigma=4, p2_amplitude=4, p2_center=3, p2_sigma=4, kinetics=4))
        params[meth, window, step] = p

print('Moving-window processing duration: %d seconds.' % (time.time() - t1))

"""
Explanation: Moving-window processing
End of explanation
"""


#moving_window_params = dict(start=0, stop=dsc.time_max, step=1, window=30)
moving_window_params

ds_slices_high = bext.moving_window_chunks(ds_high, **moving_window_params)

df = bext.moving_window_dataframe(**moving_window_params) - start_time
df['size_mean'] = [di.nt_.mean() for di in ds_slices]
df['size_max'] = [di.nt_.max() for di in ds_slices]
df['num_bursts'] = [di.num_bursts[0] for di in ds_slices]
df['burst_width'] = [di.mburst_.width.mean()*di.clk_p*1e3 for di in ds_slices]
df['burst_width_high'] = [di.mburst_.width.mean()*di.clk_p*1e3 for di in ds_slices_high]
df['phrate_mean'] = [di.max_rate_.mean() for di in ds_slices]

df = df.round(dict(tmean=1, tstart=1, tstop=1, size_mean=2, size_max=1, 
                   burst_width=2, burst_width_high=2, phrate_mean=1))
df

labels = ('num_bursts', 'burst_width', 'size_mean', 'phrate_mean',)
fig, axes = plt.subplots(len(labels), 1, figsize=(12, 3*len(labels)))

for ax, label in zip(axes, labels):
    ax.plot('tstart', label, data=df)
    ax.legend(loc='best')
    #ax.set_ylim(0)

# %%timeit -n1 -r1
# meth = 'em'
# print('    >>> Fitting method %s' % meth, flush=True)
# p = kinetics_fit(ds_slices, params_fixed, method=meth)

# %%timeit -n1 -r1
# meth = 'hist'
# print('    >>> Fitting method %s' % meth, flush=True)
# p = kinetics_fit(ds_slices, params_fixed, method=meth)

# %%timeit -n1 -r1
# meth = 'll'
# print('    >>> Fitting method %s' % meth, flush=True)
# p = kinetics_fit(ds_slices, params_fixed, method=meth)

out_fname = 'results/%s_burst_data_vs_time__window%ds_step%ds.csv' % (
    Path(filename).stem, moving_window_params['window'], moving_window_params['step'])
out_fname

df.to_csv(out_fname)

"""
Explanation: Burst-data
End of explanation
"""


# np.abs((params['em', 30, 1]  - params['ll', 30, 1]).p2_amplitude).max()

methods = ('em', 'll', 'hist')

for meth in methods:
    plt.figure(figsize=(14, 3))
    plt.plot(params['em', windows[0], step].index, params['em', windows[0], step].kinetics, 'h', color='gray', alpha=0.2)
    plt.plot(params['em', windows[1], step].index, params['em', windows[1], step].kinetics, 'h', alpha=0.3)

# (params['em', 5, 1].kinetics - params['ll', 5, 1].kinetics).plot()

for window in windows:
    for meth in methods:
        out_fname = ('results/' + Path(filename).stem + 
                     '_%sfit_ampl_only__window%ds_step%ds.csv' % (meth, window, step))
        print('- Saving: ', out_fname)
        params[meth, window, step].to_csv(out_fname)

d

"""
Explanation: Population fraction
End of explanation
"""


