#!/usr/bin/env python
# coding: utf-8

# # [PyBroMo](http://tritemio.github.io/PyBroMo/) - 1. Disk-single-core - Simulate 3D trajectories

# <small>
# *This notebook is part of [PyBroMo](http://tritemio.github.io/PyBroMo/) a 
# python-based single-molecule Brownian motion diffusion simulator 
# that simulates confocal [smFRET](http://en.wikipedia.org/wiki/Single-molecule_FRET)
# experiments. You can find the full list of notebooks in 
# [Usage Examples](http://tritemio.github.io/PyBroMo/#usage-examples).*
# </small>

# ## *Overview*
# 
# *In this notebook we show how to perform a 3-D trajectories simulation of a set of freely diffusing molecules. The simulation computes (and saves!) 3-D trajectories and emission rates due to a confocal excitation PSF for each single molecule. Depending on the simulation length, the required disk space can be significant (~ 750MB per minute of simulated diffusion).*
# 
# *For more info see [PyBroMo Homepage](http://tritemio.github.io/PyBroMo/)*.

# ## Simulation setup

# Together with a few standard python libraries we import **PyBroMo** using the short name `pbm`. 
# All **PyBroMo** functions will be available as `pbm.`*something*.

# In[ ]:


get_ipython().run_line_magic('matplotlib', 'inline')
import numpy as np
import tables
import matplotlib.pyplot as plt
import seaborn as sns
import pybromo as pbm
print('Numpy version:', np.__version__)
print('PyTables version:', tables.__version__)
print('PyBroMo version:', pbm.__version__)


# Then we define the simulation parameters:

# In[ ]:


# Initialize the random state
rs = np.random.RandomState(seed=1)
print('Initial random state:', pbm.hash_(rs.get_state()))

# Diffusion coefficient
Du = 12.0            # um^2 / s
D1 = Du*(1e-6)**2    # m^2 / s
D2 = D1/2

# Simulation box definition
box = pbm.Box(x1=-4.e-6, x2=4.e-6, y1=-4.e-6, y2=4.e-6, z1=-6e-6, z2=6e-6)

# PSF definition
psf = pbm.NumericPSF()

# Particles definition
P = pbm.Particles(num_particles=20, D=D1, box=box, rs=rs)
P.add(num_particles=15, D=D2)

# Simulation time step (seconds)
t_step = 0.5e-6

# Time duration of the simulation (seconds)
t_max = 1

# Particle simulation definition
S = pbm.ParticlesSimulation(t_step=t_step, t_max=t_max, 
                            particles=P, box=box, psf=psf)

print('Current random state:', pbm.hash_(rs.get_state()))


# The most important line is the last line which creates an object `S` 
# that contains all the simulation parameters (it also contains methods to run 
# the simulation). You can print `S` and check the current parameters:

# In[ ]:


S


# or check the required RAM for the current parameters:

# In[ ]:


S.print_sizes()


# > **NOTE:** This is the maximum in-memory array size when using a single chunk. 
# > In the following, we simulate the diffusion in smaller time windows (chunks), 
# > thus requiring only a few tens MB of RAM, regardless of the simulated duration.

# ## Brownian motion simulation
# 
# In the brownian motion simulation we keep using the same random state object `rs`. 
# Initial and final state are saved so the same simulation can be reproduced. 
# See [PyBroMo - A1. Reference - Data format and internals.ipynb](PyBroMo - A1. Reference - Data format and internals.ipynb) 
# for more info on the random state.

# In[ ]:


print('Current random state:', pbm.hash_(rs.get_state()))


# In[ ]:


#%%timeit -n1 -r1
S.simulate_diffusion(total_emission=False, save_pos=True, verbose=True, rs=rs, chunksize=2**19, chunkslice='times')


# In[ ]:


print('Current random state:', pbm.hash_(rs.get_state()))


# The normalized emission rate (peaks at 1) for each particle is stored 
# in a 2D pytable array and accessible through the `emission` attribute:

# In[ ]:


print('Simulation file size: %d MB' % (S.store.h5file.get_filesize()/1024./1024.))


# In[ ]:


S.compact_name()


# In[ ]:


S.particles.diffusion_coeff_counts


# In[ ]:


S.store.close()


# # Load trajectories

# In[ ]:


S = pbm.ParticlesSimulation.from_datafile('0168', mode='w')


# ## Plotting the emission

# In[ ]:


#from IPython.display import display


# In[ ]:


def plot_emission(S, s=0, size=2e6, slice_=None, em_th=0.01, save=False, figsize=(9, 4.5)):
    if slice_ is None:
        slice_ = (s*size, (s+1)*size)
    slice_ = slice(*slice_)
    em = S.emission[:, slice_]
    dec = 1 if slice_.step is None else slice_.step
    t_step = S.t_step*dec
    t = np.arange(em.shape[1])*(t_step*1e3)
    fig, ax = plt.subplots(figsize=figsize)
    for ip, em_ip in enumerate(em):
        if em_ip.max() < em_th: continue
        plt.plot(t, em_ip, label='P%d' % ip)
    ax.set_xlabel('Time (ms)')
    
    rs_hash = pbm.hash_(S.traj_group._v_attrs['init_random_state'])[:3]
    ax.set_title('%ds ID-EID: %d-%d, sim rs = %s, part rs = %s' %\
              (s, S.ID, S.EID, rs_hash, S.particles.rs_hash[:3]))
    ax.legend(bbox_to_anchor=(1.03, 1), loc=2, borderaxespad=0.)
    if save:
        plt.savefig('em %ds ID-EID %d-%d, rs=%s' %\
                (s, S.ID, S.EID, rs_hash), 
                dpi=200, bbox_inches='tight')
    #plt.close(fig)
    #display(fig)
    #fig.clear()


# In[ ]:


plot_emission(S, slice_=(0, 2e6, 10))


# In[ ]:


def plot_tracks(S, slice_=None, particles=None):
    if slice_ is None:
        slice_ = (0, 100e3, 100)
    duration = (slice_[1] - slice_[0])*S.t_step
    slice_ = slice(*slice_)
    
    if particles is None:
        particles = range(S.num_particles)
    
    fig, AX = plt.subplots(1, 2, figsize=(11, 5), sharey=True)
    plt.subplots_adjust(left=0.05, right=0.93, top=0.95, bottom=0.09,
                        wspace=0.05)
    plt.suptitle("Total: %.1f s, Visualized: %.2f ms" % (
                 S.t_step*S.n_samples, duration*1e3))

    for ip in particles:
        x, y, z = S.position[ip, :, slice_]
        x0, y0, z0 = S.particles[ip].r0
        plot_kwargs = dict(ls='', marker='o', mew=0, ms=2, alpha=0.5, 
                           label='P%d' % ip)
        l, = AX[0].plot(x*1e6, y*1e6, **plot_kwargs)
        AX[1].plot(z*1e6, y*1e6, color=l.get_color(), **plot_kwargs)
        #AX[1].plot([x0*1e6], [y0*1e6], 'o', color=l.get_color())
        #AX[0].plot([x0*1e6], [z0*1e6], 'o', color=l.get_color())

    AX[0].set_xlabel("x (um)")
    AX[0].set_ylabel("y (um)")
    AX[1].set_xlabel("z (um)")

    sig = np.array([0.2, 0.2, 0.6])*1e-6
    ## Draw an outline of the PSF
    a = np.arange(360)/360.*2*np.pi
    rx, ry, rz = (sig)  # draw radius at 3 sigma
    AX[0].plot((rx*np.cos(a))*1e6, (ry*np.sin(a))*1e6, lw=2, color='k')
    AX[1].plot((rz*np.cos(a))*1e6, (ry*np.sin(a))*1e6, lw=2, color='k')
    
    AX[0].set_xlim(-4, 4)
    AX[0].set_ylim(-4, 4)
    AX[1].set_xlim(-4, 4)
    
    if len(particles) <= 20:
        AX[1].legend(bbox_to_anchor=(1.01, 1), loc=2, borderaxespad=0.)


# In[ ]:


plot_tracks(S) #particles=[2, 5, 7, 22, 30])


# # Simulate timestamps

# In[ ]:


S.simulate_timestamps(max_rate=200e3, bg_rate=1e3)


# In[ ]:


S.simulate_timestamps(max_rate=300e3, bg_rate=2e3)


# In[ ]:


S.simulate_timestamps(max_rate=150e3, bg_rate=5e3)


# In[ ]:


S.timestamp_names


# In[ ]:


def plot_timetrace(S, n=-1, max_rate=200e3, bg_rate=1e3, binwidth=1e-3, scroll_gui=False):
    fig, ax = plt.subplots()
    plt.xlabel("Time (s)")
    plt.ylabel("Counts")
    times, part = S.get_timestamps(max_rate, bg_rate)
    clk_p = times._v_attrs['clk_p']
    binwidth_clk = binwidth/clk_p
    
    bins = np.arange(0, times[n] + 1, binwidth_clk)
    counts, _ = np.histogram(times[:n], bins=bins)
    
    bins *= (clk_p*1e3)
    plt.plot(bins[:-1], counts)
    plt.xlabel('Time (ms)')


# In[ ]:


plot_timetrace(S, max_rate=200e3, bg_rate=1e3)


# In[ ]:


plot_timetrace(S, max_rate=300e3, bg_rate=2e3)


# In[ ]:


plot_timetrace(S, max_rate=150e3, bg_rate=5e3)


# In[ ]:


plot_emission(S, slice_=(0, 2e6, 100))


# In[ ]:


S.get_timestamps(200e3, 1e3)


# In[ ]: