%matplotlib inline

from pprint import pprint

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pymc3 as mc
import spacepy.toolbox as tb
import spacepy.plot as spp
import tqdm
from scipy import stats
import seaborn as sns
sns.set()


%matplotlib inline

"""
Explanation: Computing dead time 2 source method
Techniques for Nuclear and Particle Physics Experiments
A How-to Approach
Authors: Leo, William R.
We measure rate1 rate2 and rate12
$n1 = \frac{R1}{1-R1\tau}, n2 = \frac{R2}{1-R2\tau}$
$n1+n2 = \frac{R12}{1-R12\tau}$
this then smplifies to:
$\tau = \frac{R1R2 - [R1R2(R12-R1)(R12-R2)]^{1/2}}{R1R2R12}$
End of explanation
"""


np.random.seed(8675309)

strength1 = 100
strength2 = 70

n_exp = 10

R1 = np.random.poisson(strength1, size=n_exp)
R2 = np.random.poisson(strength2, size=n_exp)
R12 = np.random.poisson(strength1+strength2, size=n_exp)
Rate = np.vstack((R1, R2, R12))
print(Rate)
print(Rate.shape)

print(Rate[0])
print(Rate[1])
print(Rate[2])

print(R1, R2, R12)
print(R1*R2-np.sqrt(R1*R2*(R12-R1)*(R12-R2))/(R1*R2*R12))

def tau_fn(R1, R2, R12):
    return R1*R2-np.sqrt(R1*R2*(R12-R1)*(R12-R2))/(R1*R2*R12)
print(tau_fn(R1, R2, R12))

matplotlib.pyplot.rc('figure', figsize=(10,10))
matplotlib.pyplot.rc('lines', lw=3)


plt.plot(Rate.T)


"""
Explanation: Generate some data
End of explanation
"""


with mc.Model() as model:
    mu1 = mc.Uniform('mu1', 0, 1000) # true counting rate
    mu2 = mc.Uniform('mu2', 0, 1000) # true counting rate
    mu12 = mc.Uniform('mu12', 0, 1000) # true counting rate
    R1 = mc.Poisson('R1', mu=mu1, observed=Rate[0])  # measured
    R2 = mc.Poisson('R2', mu=mu2, observed=Rate[1])  # measured
    R12 = mc.Poisson('R12', mu=mu12, observed=Rate[2])  # measured
    tau = mc.Deterministic('tau', tau_fn(R1, R2, R12))
    
    start = mc.find_MAP()
    trace = mc.sample(10000, start=start, njobs=4)

mc.summary(trace)

mc.traceplot(trace, combined=True)

tau_fn(trace['R1'][100:110], trace['R2'][100:110], trace['R12'][100:110])

"""
Explanation: So what are the errors in each measurement?
End of explanation
"""


