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

# # Table of Contents
#  <p>

# In[11]:


get_ipython().run_line_magic('matplotlib', 'inline')
import numpy as np
import statsmodels.api as sm
import pymc3 as pm
#import theano.tensor as tt
from theano import tensor as T
from matplotlib import pylab as plt
# import pystan
import seaborn as sns

np.random.seed(1)
n1 = 3000
n2 = 1500
n = n1 + n2

mu1 = 1
mu2 = 8

size = 1.2

data1 = np.random.negative_binomial(size, size/(mu1 + size), n1)
data2 = np.random.negative_binomial(size, size/(mu2 + size), n2)
data = np.concatenate([data1, data2])


# In[19]:


fig, axes = plt.subplots(2, 1, sharex=True, sharey=True)
axes[0].hist(data1)
axes[1].hist(data2);


# In[30]:


with pm.Model() as model:
    
    p = pm.Uniform( "p", 0 , 1) 
    ber = pm.Bernoulli( "ber", p = p, shape=len(data)) 

    size = pm.HalfCauchy('size', beta=2.5)

    mean = pm.Lognormal('mean', 1, 100, shape=2 )
    mu = pm.Deterministic('mu', mean[ber])
    process = pm.NegativeBinomial('obs', mu, alpha=size, observed=data)

with model:
    trace = pm.sample(100, njobs=4)


# In[34]:


pm.traceplot(trace, ['mean'])


# In[36]:


pm.summary(trace[40:], ['mean'])