# Must have R, RStudio, RTools, and JAGS installed.  I have R3.3.1, RSTudio 0.97.551, RTools 3.4, and JAGS 4.2.0.
# Must install.package(rjags)
# 
# Model
# K = number of species in sample
# N = number of possible species in the sample
# M = number of short reads
# theta[N] = Indicator of species in sampe, species are numbered 1-N; 0 out, >0 in; prior multinomial(priorspecies,K)
# alpha[N]= Abundance of species in the sample; 
# prior dirichlet with those in sample have prior proportion 1/K thse not in sample have prior proportion 1/N
# z[M] = latent, unknown assignment of reads to species; possible values are 1-N; each read is assigned a value using alpha[] probabilities
# r[M] = list of BLAST alignments associated with each reads
# r[m] ~ multinomial(length(r[m], coMatch[z[m],]); this is currently one read and a trick is used to get the likelihood
# coMatch[N, N] = given true species (row), the likelihood of aligning to another species (col)


library(rjags)
 
K <- 5
N <- 100
M <- 1000
theta <- rep(0,N)
theta[1:K] <- 1
alpha <- rep(0,N)
alpha[theta] <- 1/K

# First example has only one BLAST return that includes 
# 1 (1/50th), 
# 2 (9/50th), 
# 3 (2/50th), 
# 4 (2/50th), 
# 5 (2/50th), 
# 6 (10/50th), 
# 8 (2/50th), 
# 9 (3/50th), 
# 10 (8/50th), and 
# 50 (11/50th)
r <- sample(c(1,2,2,2,2,2,2,2,2,2,3,3,4,4,5,5,6,6,6,6,6,6,6,6,6,6,8,8,9,9,9,10,10,10,10,10,10,10,10,50,50,50,50,50,50,50,50,50,50,50), size=M, replace=T)

# Sepcify a simple coMatch matrix that makes all off-diagonal entries low probability
likelihood <- matrix(0.1,nrow=100, ncol=100)
diag(likelihood) <- rep(1,100)


write('
      model {
      # Likelihood:
      for( i in 1 : M ) {
        zeros[i] ~ dpois(phi[i])         # likelihood is exp( - phi[i])
        phi[i] <- -1*log(likelihood[r[i], z[i]])         
        z[i] ~ dcat(alpha[])  
      }
      
    # Prior:

    alpha ~ ddirch(abundance)
    for(j in 1:N){
      abundance[j] <- min(theta[j]*1/K, 1/N)
      }
    theta ~ dmulti(priorSpecies, K)

    }', "mixturemodel4.bug")

# Hyperparameters  	  
#  priorAbundance <- rep(1/N,N)
#  priorAbundance[theta] <- 1
  priorSpecies <- rep(1/N,N)
  
# false data variable for implementing a vector likelihood function in rjags
  zeros <- rep(0,M)


# Create a data list
dataList <- list(
  "r"= r,
  "zeros" = zeros,
  "likelihood" = likelihood,
  "K" = K,
  "N" = N,
  "M"= M, 
  "priorSpecies" = priorSpecies)

# List of parameters to be monitored  
parameters <- c(
  "alpha",
  "theta")

# Set initial values
initsValues <- list(
  "alpha" = alpha,
  "theta" = theta)


# JAGS Set-up
adaptSteps <- 50              #number of steps to "tune" the samplers
burnInSteps <- 50             #number of steps to "burn-in" the samplers
nChains <- 1                    #number of chains to run
numSavedSteps <- 50           #total number of steps in chains to save
thinSteps <- 1                  #number of steps to "thin" (1=keep every step)
nIter <- ceiling((numSavedSteps*thinSteps )/nChains) 	#steps per chain

# Create, initialize, and adapt the model
jagsModel <- jags.model("mixturemodel4.bug", 
                        data=dataList, 
                        inits=initsValues, 
                        n.chains=nChains, 
                        n.adapt=adaptSteps)