%matplotlib inline

import numpy as np
import scipy
import scipy.stats
from scipy.optimize import minimize
import matplotlib.pyplot as plt
from matplotlib.pyplot import hist, plot

import lmfit
from lmfit import Parameters, report_fit

time_step_s = (60e-9/4096)               # time step in seconds (S.I.)
time_step_ns = time_step_s*1e9           # time step in nano-seconds
time_nbins = 2048                        # number of time bins

time_idx = np.arange(time_nbins)         # time axis in index units
time_ns = time_idx*time_step_ns          # time axis in nano-seconds

def exgauss(x, mu, sig, tau):
    lam = 1./tau
    return 0.5*lam * np.exp(0.5*lam * (2*mu + lam*(sig**2) - 2*x)) *\
           scipy.special.erfc((mu + lam*(sig**2) - x)/(np.sqrt(2)*sig))

irf_sig = 0.033
irf_mu = 5*irf_sig
irf_lam = 0.1
x_irf = np.arange(0, (irf_lam*15 + irf_mu)/time_step_ns)
p_irf = exgauss(x_irf, irf_mu/time_step_ns, irf_sig/time_step_ns,
                irf_lam/time_step_ns)
irf = scipy.stats.rv_discrete(name='irf', values=(x_irf, p_irf))
x_irf.size

plot(x_irf, p_irf)
plt.yscale('log')
irf_mu/time_step_ns

assert np.allclose(p_irf.sum(), 1)

irf = scipy.stats.rv_discrete(name='irf', values=(x_irf, p_irf))

fig, ax = plt.subplots(1, 2, figsize=(12, 4))
ax[0].plot(x_irf, irf.pmf(x_irf));
ax[0].set_title('IRF PDF')
ax[1].hist(irf.rvs(size=200), bins=x_irf, histtype='step')
ax[1].set_title('Random samples from the IRF distribution');
ax[0].set_yscale('log')
#ax[1].set_yscale('log')

assert (p_irf == irf.pmf(x_irf)).all()

np.random.seed(2)

num_samples = 5e5
tau = 2e-9
baseline_fraction = 0.03
offset = 2e-9

sample_decay = np.random.exponential(scale=tau/time_step_s,
                                     size=num_samples) + offset/time_step_s
sample_decay += irf.rvs(size=num_samples)
sample_baseline = np.random.randint(low=0, high=time_nbins,
                                    size=baseline_fraction*num_samples)
sample_tot = np.hstack((sample_decay, sample_baseline))
decay_hist, bins = np.histogram(sample_tot, bins=np.arange(time_nbins + 1))

baseline_true = baseline_fraction*num_samples / time_nbins
baseline_true

hist_sim_decay, _, _ = hist(sample_tot, bins=np.arange(time_nbins+1), histtype='step');
plt.yscale('log')
(hist_sim_decay > 0).all()

def model(x, tau, ampl, baseline, offset, irf_pdf=irf.pmf(x_irf)):
    """
    Model function used to fit the data.
    """
    y = np.zeros(x.size)
    pos_range = (x - offset) >= 0
    y[pos_range] = ampl*np.exp(-(x[pos_range] - offset)/tau)
    z = np.convolve(y, irf_pdf, mode='same')
    z += baseline
    return z

def residuals(params, x, y, weights):
    """
    Returns the array of residuals for the current parameters.
    """
    tau = params['tau'].value
    baseline = params['baseline'].value
    offset = params['offset'].value
    ampl = params['ampl'].value
    return (y - model(x, tau, ampl, baseline, offset))*weights

def plot_fit(time, ydata, params, yscale='log', zoom_origin=False):
    """
    Function to plot data and model function.
    """
    plot(time, ydata, marker='.')
    plot(time, model(time, **{k: v.value for k, v in params.items()}),
         color='r', lw=2.5, alpha=0.7)
    if yscale == 'log':
        plt.yscale('log')
        plt.ylim(1)
    if zoom_origin:
        dt = time[1] - time[0]
        t0 = params['offset'] - 50*dt
        plt.xlim(t0 - 50*dt, t0 + 50*dt)


def loglike(params, x, ydata):
    tau, baseline, offset, ampl = params
    ymodel = model(x, tau, ampl, baseline, offset)
    return (ymodel - ydata*np.log(ymodel)).sum()

def from_params(params):
    return [params[k].value for k in ('tau', 'baseline', 'offset', 'ampl')]

def to_params(p, params):
    for v, k in zip(p, ('tau', 'baseline', 'offset', 'ampl')):
        params[k].value = v
    return params

weights = 1./np.sqrt(decay_hist)

weights[decay_hist == 0] = 1./np.sqrt(baseline_true)
(decay_hist == 0).any()

params = Parameters()
params.add('tau', value=2.8)
params.add('baseline', value=5)
params.add('offset', value=2.814, vary=False)
params.add('ampl', value=700)

plot_fit(time_ns, decay_hist, params)

fit_res = lmfit.minimize(residuals, params, args=(time_ns, decay_hist, weights), method='nelder')
report_fit(fit_res.params)
print 'Reduced Chi-square: %.3f' % fit_res.redchi

fit_res = lmfit.minimize(residuals, params, args=(time_ns, decay_hist, weights), method='leastsq')
report_fit(fit_res.params)
print 'Reduced Chi-square: %.3f' % fit_res.redchi

plot_fit(time_ns, decay_hist, params, zoom_origin=False)

ci = lmfit.conf_interval(fit_res)

lmfit.printfuncs.report_ci(ci)

res = minimize(loglike, x0=from_params(params), args=(time_ns, decay_hist),
               method='Nelder-Mead')

params = to_params(res.x, params)

report_fit(params)

plot_fit(time_ns, decay_hist, params, zoom_origin=False)



from IPython.core.display import HTML
HTML(open("./styles/custom2.css", "r").read())