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

# # [PyBroMo](http://tritemio.github.io/PyBroMo/) 4. Two-state dynamics - Static smFRET simulation

# <small><i>
# This notebook is part of <a href="http://tritemio.github.io/PyBroMo" target="_blank">PyBroMo</a> a 
# python-based single-molecule Brownian motion diffusion simulator 
# that simulates confocal smFRET
# experiments.
# </i></small>

# ## *Overview*
# 
# *In this notebook we generate single-polulation static smFRET data files 
# from the same diffusion trajectories. These files are needed by the [next notebook](PyBroMo - 5. Two-state dynamics - Dynamic smFRET simulation)
# to generate smFRET data with 2-state dynamics.*

# ## Loading the software
# 
# Import all the relevant libraries:

# In[ ]:


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


# In[ ]:


SIM_PATH = 'data/'


# # Define populations

# We assume a $\gamma = 0.7$ and two populations, one with $E_{PR}=0.75$ and the other $E_{PR}=0.4$

# In[ ]:


Epr1 = 0.75
Epr2 = 0.4


# The corrected $E$ for the two populations are:

# In[ ]:


gamma = 0.7
E1 = Epr1 /(Epr1 * (1 - gamma) + gamma)
E2 = Epr2 /(Epr2 * (1 - gamma) + gamma)
E1, E2


# The simulation takes the uncorrected $E_{PR}$ as input. 
# 
# We want to simulate a second population that has measured brightness scaling as if the difference was
# only due to the $\gamma$ factor. Using the definitions $\Lambda$ and $\Lambda_\gamma$ from 
# [(Ingargiola 2017)](https://doi.org/10.1101/156182),
# we can use the relation:
# 
# $$\frac{E}{E_{PR}} = \frac{\Lambda}{\Lambda_\gamma}$$
# 
# Solving for $\Lambda_\gamma$ or $\Lambda$ we get:
# 
# $$ \Lambda_\gamma  = \Lambda\frac{E_{PR}}{E}$$
# 
# $${\Lambda} = {\Lambda_\gamma}\frac{E}{E_{PR}} $$
# 
# Since $\Lambda_\gamma$ is gamma corrected does not depend on $E$. We can compute it from 
# the parameters of population 1 and then use it for finding $\Lambda$ for population 2:

# In[ ]:


Λ1 = 200e3  # kcps, the detected (i.e. uncorrected) peak emission rate for population 1
Λγ = Λ1 * Epr1 / E1
Λ2 = Λγ * E2 / Epr2
Λ2


# In[ ]:


Λ2 = np.round(Λ2, -3)
Λ2


# # Create smFRET data-files

# ## Create a file for storing timestamps
# 
# Here we load a diffusion simulation opening a file to save
# timstamps in *write* mode. Use `'a'` (i.e. *append*) to keep 
# previously simulated timestamps for the given diffusion.

# In[ ]:


S = pbm.ParticlesSimulation.from_datafile('0eb9', mode='a', path=SIM_PATH)


# In[ ]:


S.particles.diffusion_coeff_counts


# In[ ]:


#S = pbm.ParticlesSimulation.from_datafile('44dc', mode='a', path=SIM_PATH)


# ## Simulate timestamps of smFRET
# 
# We want to simulate two separate smFRET files representing two static populations.
# 
# We start definint the simulation parameters for population 1 with the following syntax:

# In[ ]:


params1 = dict(
    em_rates = (Λ1,),       # Peak emission rates (cps) for each population (D+A)
    E_values = (Epr1,),     # FRET efficiency for each population
    num_particles = (35,),  # Number of particles in each population
    bg_rate_d = 900,        # Poisson background rate (cps) Donor channel
    bg_rate_a = 600,        # Poisson background rate (cps) Acceptor channel
    )


# We can now define population 2:

# In[ ]:


params2 = dict(
    em_rates = (Λ2,),       # Peak emission rates (cps) for each population (D+A)
    E_values = (Epr2,),     # FRET efficiency for each population
    num_particles = (35,),  # Number of particles in each population
    bg_rate_d = 900,        # Poisson background rate (cps) Donor channel
    bg_rate_a = 600,        # Poisson background rate (cps) Acceptor channel
    )


# Finally, we also define a static mixture of the two populations:

# In[ ]:


params_mix = dict(
    em_rates = (Λ1, Λ2),       # Peak emission rates (cps) for each population (D+A)
    E_values = (Epr1, Epr2),   # FRET efficiency for each population
    num_particles = (20, 15),  # Number of particles in each population
    bg_rate_d = 900,           # Poisson background rate (cps) Donor channel
    bg_rate_a = 600,           # Poisson background rate (cps) Acceptor channel
    )


# ### Simulate static population 1

# **Population 1**: Create the object that will run the simulation and print a summary:

# In[ ]:


mix_sim = pbm.TimestapSimulation(S, **params1)
mix_sim.summarize()


# Run the simualtion:

# In[ ]:


rs = np.random.RandomState(1234)
mix_sim.run(rs=rs, overwrite=False, skip_existing=True)


# Save simulation to a smFRET [Photon-HDF5](http://photon-hdf5.org) file:

# In[ ]:


mix_sim.save_photon_hdf5(identity=dict(author='Antonino Ingargiola', 
                                       author_affiliation='UCLA'))


# ### Simulate static population 2

# **Population 2**: Create the object that will run the simulation and print a summary:

# In[ ]:


mix_sim = pbm.TimestapSimulation(S, **params2)
mix_sim.summarize()


# In[ ]:


rs = np.random.RandomState(1234)
mix_sim.run(rs=rs, overwrite=False, skip_existing=True)


# In[ ]:


mix_sim.save_photon_hdf5(identity=dict(author='Antonino Ingargiola', 
                                       author_affiliation='UCLA'))


# ### Simulate static mixture

# **Static mixture**: Create the object that will run the simulation and print a summary:

# In[ ]:


mix_sim = pbm.TimestapSimulation(S, **params_mix)
mix_sim.summarize()


# In[ ]:


rs = np.random.RandomState(1234)
mix_sim.run(rs=rs, overwrite=False, skip_existing=True)


# In[ ]:


mix_sim.save_photon_hdf5(identity=dict(author='Antonino Ingargiola', 
                                       author_affiliation='UCLA'))


# In[ ]:


get_ipython().system("rsync -av --exclude 'pybromo_*.hdf5' /mnt/archive/Antonio/pybromo /mnt/wAntonio/dd")


# # Burst analysis
# 
# The generated Photon-HDF5 files can be analyzed by any smFRET burst
# analysis program. Here we show an example using the opensource
# [FRETBursts](https://github.com/tritemio/FRETBursts/) program:

# In[ ]:


import fretbursts as fb


# In[ ]:


filepath = list(Path(SIM_PATH).glob('smFRET_*'))
filepath


# In[ ]:


d = fb.loader.photon_hdf5(str(filepath[0]))


# In[ ]:


d


# In[ ]:


d.A_em


# In[ ]:


fb.dplot(d, fb.timetrace);


# In[ ]:


d.calc_bg(fun=fb.bg.exp_fit, tail_min_us='auto', F_bg=1.7, time_s=5)


# In[ ]:


fb.dplot(d, fb.timetrace_bg)


# In[ ]:


d.burst_search(F=7)


# In[ ]:


d.num_bursts


# In[ ]:


ds = d.select_bursts(fb.select_bursts.size, th1=20)


# In[ ]:


ds.num_bursts


# In[ ]:


with plt.rc_context({#'font.size': 10, 
                     #'savefig.dpi': 200, 
                     'figure.dpi': 150}):
    for i in range(3):
        fig, ax = plt.subplots(figsize=(100, 3))
        fb.dplot(d, fb.timetrace, binwidth=0.5e-3, tmin=i*10, tmax=(i+1)*10, bursts=True, 
                 plot_style=dict(lw=1), ax=ax);
        ax.set_xlim(i*10, (i+1)*10);
        display(fig)
        plt.close(fig)


# In[ ]:


fb.dplot(ds, fb.hist_fret, pdf=False)
plt.axvline(0.4, color='k', ls='--');


# In[ ]:


fb.bext.burst_data(ds)


# In[ ]:


fb.dplot(d, fb.hist_size)


# In[ ]:


fb.dplot(d, fb.hist_width)