from fretbursts import *

sns = init_notebook(apionly=True)
print('seaborn version: ', sns.__version__)

# Tweak here matplotlib style
import matplotlib as mpl
mpl.rcParams['font.sans-serif'].insert(0, 'Arial')
mpl.rcParams['font.size'] = 12
%config InlineBackend.figure_format = 'retina'

"""
Explanation: Example - Customize the μs-ALEX histogram
This notebook is part of smFRET burst analysis software FRETBursts.

In this notebook shows how to plot different styles of μs-ALEX histograms and $E$ and $S$ marginal distributions.
For a complete tutorial on burst analysis see 
FRETBursts - us-ALEX smFRET burst analysis.
End of explanation
"""


url = 'http://files.figshare.com/2182601/0023uLRpitc_NTP_20dT_0.5GndCl.hdf5'
download_file(url, save_dir='./data')
full_fname = "./data/0023uLRpitc_NTP_20dT_0.5GndCl.hdf5"

d = loader.photon_hdf5(full_fname)
loader.alex_apply_period(d)
d.calc_bg(bg.exp_fit, time_s=1000, tail_min_us=(800, 4000, 1500, 1000, 3000))
d.burst_search(L=10, m=10, F=6)
ds = d.select_bursts(select_bursts.size, add_naa=True, th1=30)

"""
Explanation: Get and process data
End of explanation
"""


alex_jointplot(ds)

"""
Explanation: ALEX joint plot
The alex_jointplot function allows plotting an ALEX histogram with marginals.
This is how it looks by default:
End of explanation
"""


alex_jointplot(ds, vmax_fret=False)

alex_jointplot(ds, vmax_fret=False, marginal_color=8)

alex_jointplot(ds, vmax_fret=False, marginal_color=7)

alex_jointplot(ds, kind='kde')

"""
Explanation: The inner plot in an hexbin plot, basically a 2D histogram with hexagonal bins.
This kind of histograms resembles a scatter plot when sample size is small,
and is immune from grid artifacts typical of rectangular grids.
For more info for hexbin see this document.
The marginal plots are histograms with an overlay KDE plot. 
The same FRETBursts function that plots standalone E and S histograms 
is used here to plot the marginals in the joint plot.
Below I show how to customize appearance and type of this plot.
Changing colors
By default the colormap range is computed on the range S=[0.2, 0.8],
so that the FRET populations (S ~ 0.5) have more contrast.
To normalize the colormap to the whole data use the vmax argument:
End of explanation
"""


alex_jointplot(ds, vmax=40)

"""
Explanation: Or you can manually choose the max value mapped by the colormap (vmax):
End of explanation
"""


alex_jointplot(ds, cmap='plasma')

"""
Explanation: Changing the colormap will affect both inner and marginal plots:
End of explanation
"""


alex_jointplot(ds, cmap='plasma', marginal_color=83)

"""
Explanation: To pick a different color from the colormap for the marginal histograms use histcolor_id:
End of explanation
"""


alex_jointplot(ds, kind='scatter')

alex_jointplot(ds, kind='kde')

dsf = ds.select_bursts(select_bursts.naa, th1=40)
alex_jointplot(dsf, kind='kde',
               joint_kws={'shade': False, 'n_levels': 12, 'bw': 0.04})

"""
Explanation: Kinds of joint-plots
The inner plot can be changed to a scatter plot or a KDE plot:
End of explanation
"""


plt.figure(figsize=(5,5))
hexbin_alex(ds)

"""
Explanation: No marginals
Finally, we can plot only the hexbin 2D histogram without marginals:
End of explanation
"""


g = alex_jointplot(ds)
g.ax_marg_x.grid(False)
g.ax_marg_y.grid(False)
g.ax_joint.set_xlim(-0.1, 1.1)
g.ax_joint.set_ylim(-0.1, 1.1)

"""
Explanation: Figure layout
You can get an handle of the different axes in the figure for layout customization:
End of explanation
"""


g = alex_jointplot(ds)
g.ax_marg_x.grid(False)
g.ax_marg_y.grid(False)
g.ax_joint.set_xlim(-0.19, 1.19)
g.ax_joint.set_ylim(-0.19, 1.19)
plt.subplots_adjust(wspace=0, hspace=0)
g.ax_marg_y.spines['bottom'].set_visible(True)
g.ax_marg_x.spines['left'].set_visible(True)
g.ax_marg_y.tick_params(reset=True, bottom=True, top=False, right=False, labelleft=False)
g.ax_marg_x.tick_params(reset=True, left=True, top=False, right=False, labelbottom=False)

g = alex_jointplot(ds)
g.ax_marg_x.grid(False)
g.ax_marg_y.grid(False)
g.ax_joint.set_xlim(-0.19, 1.19)
g.ax_joint.set_ylim(-0.19, 1.19)
plt.subplots_adjust(wspace=0, hspace=0)
g.ax_marg_y.tick_params(reset=True, bottom=True, top=False, right=False, labelleft=False)
g.ax_marg_x.tick_params(reset=True, left=True, top=False, right=False, labelbottom=False)

g = alex_jointplot(ds)
g.ax_marg_x.grid(False, axis='x')
g.ax_marg_y.grid(False, axis='y')
g.ax_joint.set_xlim(-0.19, 1.19)
g.ax_joint.set_ylim(-0.19, 1.19)
plt.subplots_adjust(wspace=0, hspace=0)

"""
Explanation: alex_jointplot returns g which contains the axis handles (g.ax_join, g.ax_marg_x, g.ax_marg_y).
The object g is a seaborn.JointGrid.
End of explanation
"""


alex_jointplot(ds, marginal_kws={'show_kde': False})

"""
Explanation: Arguments of inner plots
Additional arguments can be passed to the inner or marginal plots passing 
a dictionary to joint_kws and marginal_kws respectively.
The marginal plots are created by hist_burst_data 
which is the same function used to plot standalone E and S histograms
in FRETBursts. 
For example, we can remove the KDE overlay like this:
End of explanation
"""


from ipywidgets import widgets, interact, interactive, fixed
from IPython.display import display, display_png, display_svg, clear_output
from IPython.core.pylabtools import print_figure

cmaps = ['viridis', 'plasma', 'inferno', 'magma',
         'afmhot', 'Blues', 'BuGn', 'BuPu', 'GnBu', 'YlGnBu',
         'coolwarm', 'RdYlBu', 'RdYlGn', 'Spectral',]# 'icefire']  uncomment if using seaborn 0.8

@interact(overlay = widgets.RadioButtons(options=['fit model', 'KDE'], value='KDE'),
          binwidth = widgets.FloatText(value=0.03, min=0.01, max=1),
          bandwidth = widgets.FloatText(value=0.03, min=0.01, max=1),
          gridsize = (10, 100),
          min_size=(10, 500, 5),
          cmap=widgets.Dropdown(value='Spectral', options=cmaps),
          reverse_cmap = True,
          vmax_fret = True,
          )
def plot_(min_size=50, overlay='KDE', binwidth=0.03, bandwidth=0.03, 
          gridsize=50, cmap='Spectral', reverse_cmap=False, 
          vmax_fret=True):
    dx = d.select_bursts(select_bursts.size, add_naa=True, th1=min_size)
    bext.bursts_fitter(dx, 'E', binwidth=binwidth, bandwidth=bandwidth, 
                       model=mfit.factory_three_gaussians())
    bext.bursts_fitter(dx, 'S', binwidth=binwidth, bandwidth=bandwidth, 
                       model=mfit.factory_two_gaussians())    
    
    if reverse_cmap: cmap += '_r'

    if binwidth < 0.01: binwidth = 0.01
    if bandwidth < 0.01: bandwidth = 0.01
    if overlay == 'fit model':
        marginal_kws = dict(binwidth=binwidth, show_model=True, pdf=True, 
                            show_kde=False)
    else:
        marginal_kws = dict(binwidth=binwidth, show_kde=True, 
                            bandwidth=bandwidth)
    alex_jointplot(dx, cmap=cmap, gridsize=gridsize, vmax_fret=vmax_fret, 
                   marginal_kws=marginal_kws,)
    
    fig = gcf()
    plt.close()
    display(fig)

"""
Explanation: Interactive plot
End of explanation
"""


