# This code takes a vector of data y and estimates 
# MAC model parameters. Vector y contains the number of 
# correct responses for each participant; vector N contains 
# the total number of observations for each participant.

rm(list=ls()) # Clear workspace

#-----Helper functions------------------

# Function to sample from a TN(mu,sigma^2) truncated above 
# at value trunc.at
rtnorm.lower = function(n=1,mu=0,sigma=1,trunc.at=0)
{
  sigma * qnorm(log(runif(n)) +
    pnorm((trunc.at-mu)/sigma,log.p=T),log.p=T) + mu
}

# Function to sample from a TN(mu,sigma^2) truncated below 
# at value trunc.at 
rtnorm.upper = function(n=1,mu=0,sigma=1,trunc.at=0)
{
  -sigma * qnorm(log(runif(n)) +
    pnorm((-trunc.at+mu)/sigma,log.p=T),log.p=T) + mu
}

# Create a general rtnorm function; this function will 
# generate n normal deviates with mean mu and standard
# deviation sigma, truncated at trunc.at. n.lower values
# will be from normal truncated above 0, and n-n.lower
# from the normal truncated below 0.
rtnorm = function(n=1,mu=0,sigma=1,trunc.at=0,n.lower=1)
{
  c(rtnorm.lower(n.lower,mu,sigma,trunc.at),
  rtnorm.upper(n-n.lower,mu,sigma,trunc.at))
}

# Function to sample from an TIG(t,a,b)
rtruncinvgamma = function(n,t,a,b)
{
  u = runif(n)
  q = 1-pgamma(1/t,a,b)
  1/qgamma(1-u*q,a,b)
}


#-------End Helper Functions-------------

# Data; 27 subjects from the experiment described in Rouder,
# Morey, Speckman, and Pratte's 'Differentiating Chance 
# Performance from above-chance performance: A solution to
# the null sensitivity problem in evaluating subliminal 
# priming.'
# y = Number or correct responses
y = c(150,142,154,155,136,138,211,140,148,159,164,150,158,138, 
      148,146,163,145,180,155,148,147,134,134,167,149,147)
# N = Number of observations
N = c(284,288,287,288,288,288,288,288,285,287,288,288,288,288, 
      288,288,288,288,288,288,287,287,288,286,288,288,288)

# For housekeeping
obs.fewer=max(N)-N # How many observations LESS does each have?

# Initialize variables
M      = 50000      # Length of MCMC chain      
I      = length(y)  # Number of participants
burnin = 5000       # Throw away this many at start of chain 


w       = matrix(nrow=max(N),ncol=I)
x       = matrix(nrow=M,ncol=I)
mu      = matrix(nrow=M)
sigma.2 = matrix(nrow=M)

# Initialize prior parameters
mu_0      = 0
sigma_0.2 = 1
a_0       = -.5
b_0       = 0
t_sig2    = 1

# Starting values
p.hat = (y+1)/(N+2) # Estimate p from data for starting values
x[1,]      = qnorm(p.hat)
mu[1]      = 0
sigma.2[1] = 1
for(i in 1:I) # Find reasonable starting values for w_ij
{ 
  w[,i] = c(rtnorm(n=N[i],mu=x[1,i],sigma=1,
           trunc.at=0,n.lower=N[i]-y[i]),
           rep(NA,obs.fewer[i]))
}

# Main MCMC loop
for(m in 2:M) # Loop through iterations
{     
  for(i in 1:I)  # Loop through participants
  {  
    # Sampling x_i
    lambda_i.2 = (N[i] + (1/sigma.2[m-1]))^-1
    nu_i     = lambda_i.2*(sum(w[,i],na.rm=T) +
               mu[m-1]/sigma.2[m-1])
    gamma_i  = log(sqrt(sigma.2[m-1])) + 
               pnorm(-mu[m-1]/sqrt(sigma.2[m-1]),log.p=T) -
               log(sqrt(lambda_i.2)) - 
               pnorm(nu_i/sqrt(lambda_i.2),log.p=T) +
               (1/2)*(mu[m-1]^2/sigma.2[m-1] - nu_i^2/lambda_i.2) 
    rho_i    = 1/(1+exp(-gamma_i))
    r_i      = rbinom(1,1,rho_i) # Sample Bernoulli
    if(r_i==0){
      x[m,i] = rtnorm.upper(n=1,mu=nu_i,
                 sigma=sqrt(lambda_i.2),trunc.at=0)
    }else{
      x[m,i] = rtnorm.lower(n=1,mu=mu[m-1],
                 sigma=sqrt(sigma.2[m-1]),trunc.at=0)
    }
    # Sampling w_ij
    w[,i] = c(rtnorm(n=N[i],mu=max(0,x[m,i]),
              sigma=1,trunc.at=0,n.lower=N[i]-y[i]),
              rep(NA,obs.fewer[i]))
  } # End loop through participants
  
  # Sampling mu
  tau.2 = (I/sigma.2[m-1] + 1/sigma_0.2)^-1
  eta   = tau.2*(sum(x[m,])/sigma.2[m-1] + mu_0/sigma_0.2)
  mu[m] = rnorm(1,eta,sqrt(tau.2))
  
  # Sampling sigma^2
  sigma.2[m] = rtruncinvgamma(1, t_sig2, a_0 + I/2,
                 b_0 + sum((x[m,]-mu[m])^2)/2)

  # Report progress
  if(!(m%%500)) { cat('Iteration number:',m,'\n') }

} # End main MCMC loop

# Estimate parameters (be sure to check for convergence)
mu.estimate      = mean(mu[burnin:M]) 
sigma.2.estimate = mean(sigma.2[burnin:M])
pr.x.lt.0        = colMeans(x[burnin:M,]<0)

# Create a plot of the output
plot(y/N,pr.x.lt.0)
abline(h=.95)