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

# Original BUGS model from Link and Barker 2009
# 
# ```
# # Data 
# list(T=2, Cells=4, f=c(NA,22,363,174), 
# omegas=structure(.Data=c(0,0, 
# 1,0, 
# 0,1, 
# 1,1),.Dim=c(4,2))) 
# 
# model 
# { 
#     for (t in 1:T) 
#     { 
#         p[t] ~ dunif(0,1) 
#     }     
#     for (j in 1:Cells) 
#     { 
#         for (t in 1:T) 
#         { 
#             c[j,t] <- pow(p[t],omegas[j,t]) * pow(1-p[t],1-omegas[j,t]) 
#         } 
#         pi[j] <- prod(c[j,1:T]) 
#     } 
#     for (j in 2:Cells) 
#     { 
#         pi.obs[j] <- pi[j]/(1-pi[1]) 
#     } 
#     n <- sum(f[2:Cells]) 
#     f[2:Cells] ~ dmulti(pi.obs[2:Cells],n) 
#     pn <- 1-pi[1] 
#     n ~ dbin(pn,N) 
#     cN ~ dunif(n,10000) 
#     N <- round(cN) 
# } 
# 
# ```

# In[21]:


get_ipython().run_line_magic('matplotlib', 'inline')
import numpy as np
import pymc3 as pm
import theano.tensor as tt


# In[14]:


# Number of observers
T = 2
# Observations by neither, both, first and second observer
f = np.array([np.nan, 22, 363, 174])
# Total observations
n = np.nansum(f)
# Observer indicators
omega = np.reshape((0,0, 
                    1,0, 
                    0,1, 
                    1,1), (4, 2))


# In[26]:


with pm.Model() as Mt2:
    
    # Individual detection probabilities
    p = pm.Uniform('p', 0, 1, shape=T)
    
    # Joint probabilities
    c = p**omega * (1 - p)**(1-omega)
    π = pm.Deterministic('π', tt.prod(c, 1))
    
    # Conditional (on observation) probabilities
    π_obs = π[1:] / (1 - π[0])
    
    f_obs = pm.Potential('f_obs', pm.Multinomial.dist(n=n, p=π_obs).logp(f[1:]))
    
    # Latent population size
    N = pm.Uniform('N', 0, 10000)
    
    # Probability of being observed by one or more observers
    n_obs = pm.Potential('n_obs', pm.Binomial.dist(n=tt.floor(N), p=1-π[0]).logp(n))


# In[30]:


with Mt2:
    trace = pm.sample(1000, tune=1000)


# In[31]:


pm.plot_posterior(trace, varnames=['N'])


# In[35]:


pm.summary(trace, varnames=['p'])


# In[33]:


pm.forestplot(trace, varnames=['π'])