from fretbursts import *
sns = init_notebook()

import os
import pandas as pd
from IPython.display import display, Math

import lmfit
print('lmfit version:', lmfit.__version__)

figure_size = (5, 4)
default_figure = lambda: plt.subplots(figsize=figure_size)
save_figures = True

def savefig(filename, **kwargs):
    if not save_figures:
        return
    import os
    dir_ = 'figures/'
    kwargs_ = dict(dpi=300, bbox_inches='tight')
                   #frameon=True, facecolor='white', transparent=False)
    kwargs_.update(kwargs)
    plt.savefig(dir_ + filename, **kwargs_)
    print('Saved: %s' % (dir_ + filename))

"""
Explanation: Multi-spot vs usALEX FRET histogram comparison
Load FRETBursts software
End of explanation
"""


PLOT_DIR = './figure/'

import matplotlib as mpl
from cycler import cycler

bmap = sns.color_palette("Set1", 9)
colors = np.array(bmap)[(1,0,2,3,4,8,6,7), :]
mpl.rcParams['axes.prop_cycle'] = cycler('color', colors)

"""
Explanation: 8-spot paper plot style
End of explanation
"""


data_dir = './data/multispot/'

"""
Explanation: Data files
Data folder:
End of explanation
"""


data_dir = os.path.abspath(data_dir) + '/'
assert os.path.exists(data_dir), "Path '%s' does not exist." % data_dir

"""
Explanation: Check that the folder exists:
End of explanation
"""


from glob import glob
file_list = sorted(glob(data_dir + '*_?.hdf5'))

labels = ['12d', '7d', '17d', '22d', '27d', 'DO']
files_dict = {lab: fname for lab, fname in zip(sorted(labels), file_list)}
files_dict

"""
Explanation: List of data files in data_dir:
End of explanation
"""


_fname = 'results/Multi-spot - leakage coefficient KDE wmean DexDem.csv'
leakageM = np.loadtxt(_fname, ndmin=1)

print('Leakage coefficient:', leakageM)

"""
Explanation: Correction parameters
Multispot
Load the multispot leakage coefficient from disk (computed in Multi-spot 5-Samples analyis - Leakage coefficient fit):
End of explanation
"""


_fname = 'results/usALEX - direct excitation coefficient dir_ex_t beta.csv'
dir_ex_tM = np.loadtxt(_fname, ndmin=1)

print('Direct excitation coefficient (dir_ex_t):', dir_ex_tM)

"""
Explanation: Load the multispot direct excitation coefficient ($d_{dirT}$) from disk (computed in usALEX - Corrections - Direct excitation physical parameter):
End of explanation
"""


_fname = 'results/Multi-spot - gamma factor.csv'
gammaM = np.loadtxt(_fname, ndmin=1)

print('Multispot gamma coefficient:', gammaM)

"""
Explanation: Load the multispot gamma ($\gamma_M$) coefficient (computed in Multi-spot Gamma Fitting):
End of explanation
"""


_fname = 'results/usALEX - leakage coefficient DexDem.csv'
leakageA = np.loadtxt(_fname)

print('usALEX Leakage coefficient:', leakageA)

"""
Explanation: usALEX
Load the usALEX leakage coefficient from disk (computed in usALEX - Corrections - Leakage fit):
End of explanation
"""


_fname = 'results/usALEX - gamma factor - all-ph.csv'
gammaA = np.loadtxt(_fname)

print('usALEX Gamma-factor:', gammaA)

"""
Explanation: Load the usALEX gamma coefficient (computed in usALEX - Corrections - Gamma factor fit):
End of explanation
"""


_fname = 'results/usALEX - beta factor - all-ph.csv'
betaA = np.loadtxt(_fname)

print('usALEX Gamma-factor:', betaA)

"""
Explanation: Load the usALEX beta coefficient (computed in usALEX - Corrections - Gamma factor fit):
End of explanation
"""


_fname = 'results/usALEX - direct excitation coefficient dir_ex_aa.csv'
dir_ex_aa = np.loadtxt(_fname)

print('Direct excitation coefficient (dir_ex_aa):', dir_ex_aa)

"""
Explanation: Load the usALEX direct-excitation coefficient ($d_{exAA}$) (computed in usALEX - Corrections - Direct excitation fit):
End of explanation
"""


dir_ex_tA = betaA * dir_ex_aa
dir_ex_tA

"""
Explanation: Compute usALEX direct-excitation coefficient ($d_{exT}$) (see usALEX - Corrections - Direct excitation physical parameter):
End of explanation
"""


donor_ref = False    # False -> gamma correction is: g*nd + na
                     # True  -> gamma correction is: nd + na/g

hist_weights = 'size'

## Background fit parameters
bg_kwargs_auto = dict(fun=bg.exp_fit,
                 time_s = 30,
                 tail_min_us = 'auto',
                 F_bg=1.7,
                 )

## Burst search
F=6
dither = False
size_th = 30    # Burst size threshold (selection on corrected burst sizes)

## FRET fit parameters
bandwidth = 0.03        # KDE bandwidth
E_range = {'7d':  (0.7, 1.0), '12d': (0.4, 0.8), '17d': (0.2, 0.4), 
           '22d': (0.0, 0.1), '27d': (0.0, 0.1), 'DO': (0.0, 0.1)}
E_axis_kde = np.arange(-0.2, 1.2, 0.0002)

"""
Explanation: Parameters
Analysis parameters:
End of explanation
"""


def print_fit_report(E_pr, gamma=1, leakage=0, dir_ex_t=0, math=True):
    """Print fit and standard deviation for both corrected and uncorrected E
    Returns d.E_fit.
    """
    E_corr = fretmath.correct_E_gamma_leak_dir(E_pr, gamma=gamma, leakage=leakage, dir_ex_t=dir_ex_t)
    
    E_pr_mean = E_pr.mean()*100
    E_pr_delta = (E_pr.max() - E_pr.min())*100
    
    E_corr_mean = E_corr.mean()*100
    E_corr_delta = (E_corr.max() - E_corr.min())*100
    if math:
        display(Math(r'\text{Pre}\;\gamma\quad\langle{E}_{fit}\rangle = %.1f\%% \qquad'
                     '\Delta E_{fit} = %.2f \%%' % \
                     (E_pr_mean, E_pr_delta)))
        display(Math(r'\text{Post}\;\gamma\quad\langle{E}_{fit}\rangle = %.1f\%% \qquad'
                     '\Delta E_{fit} = %.2f \%%' % \
                     (E_corr_mean, E_corr_delta)))
    else:
        print('Pre-gamma  E (delta, mean):  %.2f  %.2f' % (E_pr_mean, E_pr_delta))
        print('Post-gamma E (delta, mean):  %.2f  %.2f' % (E_corr_mean, E_corr_delta))

"""
Explanation: Utility functions
End of explanation
"""


d = loader.photon_hdf5(files_dict[data_id])
d.calc_bg(**bg_kwargs_auto)
d.burst_search(m=10, F=F, dither=dither)

d.time_max

ds = Sel(d, select_bursts.size, th1=30, gamma=gammaM, donor_ref=donor_ref)

ds.num_bursts

# fitter = bext.bursts_fitter(ds)
# fitter.histogram(bins=np.r_[-0.2 : 1.2 : bandwidth])
# fitter.model = mfit.factory_two_gaussians(add_bridge=False, p2_center=0.4)
# fitter.fit_histogram()
# display(fitter.params['p2_center'])
# print_fit_report(fitter.params['p2_center'], gamma=gammaM, leakage=leakageM, dir_ex_t=dir_ex_tM)

dplot(ds, hist_fret);
      #show_model=True, show_fit_stats=True, fit_from='p2_center', show_fit_value=True);

d_all = ds.collapse()

d_all_chunk = Sel(d_all, select_bursts.time, time_s2=600/8)

dplot(d_all_chunk, hist_fret)

Eraw = d_all_chunk.E[0]

E = fretmath.correct_E_gamma_leak_dir(Eraw, gamma=gammaM, leakage=leakageM, dir_ex_t=dir_ex_tM)

sns.set_style('whitegrid')
%config InlineBackend.figure_format='retina'  # for hi-dpi displays

plt.hist(E, bins=np.arange(-0.2, 1.2, 0.025) + 0.5*0.025);

"""
Explanation: Multispot analysis
End of explanation
"""


bursts_usalex = pd.read_csv('results/bursts_usALEX_{sample}_{ph_sel}_F{F:.1f}_m{m}_size{th}.csv'
                            .format(sample=data_id, ph_sel='Dex', m=10, th=30, F=7), index_col=0)
bursts_usalex

Eraw_alex = bursts_usalex.E

E_alex = fretmath.correct_E_gamma_leak_dir(Eraw_alex, gamma=gammaA, leakage=leakageA, dir_ex_t=dir_ex_tA)

kws = dict(bins=np.arange(-0.2, 1.2, 0.025) + 0.5*0.025, histtype='step', lw=1.8)
plt.hist(E, label='Multispot', **kws)
plt.hist(E_alex, label='μs-ALEX', **kws)
plt.legend(loc=2)
plt.title('Sample %s: Multispot vs μs-ALEX comparison' % data_id)
plt.xlabel('FRET Efficiency')
plt.ylabel('# Bursts');
savefig('Multispot vs usALEX FRET hist comp sample %s' % data_id)

kws = dict(bins=np.arange(-0.2, 1.2, 0.025) + 0.5*0.025, histtype='step', lw=1.8, normed=True)
plt.hist(E, label='Multispot', **kws)
plt.hist(E_alex, label='μs-ALEX', **kws)
plt.legend(loc=2)
plt.title('Sample %s: Multispot vs μs-ALEX comparison' % data_id)
plt.xlabel('FRET Efficiency')
plt.ylabel('Probabiltity');
savefig('Multispot vs usALEX FRET hist comp sample %s normed' % data_id)

"""
Explanation: Comparison with usALEX
End of explanation
"""


