rm(list=ls(all=TRUE))  # clear all variables
graphics.off()  # clear all graphics
# Zenner cards
# Greg Francis
# PSY 646
# 15 February 2024

# load full data file
ZCdata<-read.csv(file="ZennerCards.csv",header=TRUE,stringsAsFactors=FALSE)

# Make data frame that will work for Stan
ZCdataClean<- data.frame(Score= ZCdata$Score, Participant= ZCdata$Participant)

# load the rethinking library
library(rethinking)
	

# Null model (p=0.2)
ZCmodel0 <- ulam(
			alist( Score ~ dbinom(1, p), 
			logit(p) <- a,
			a ~ dnorm(logit(0.2), 0.001)
			), data= ZCdataClean, log_lik=TRUE, chains=4  ) 
		
print(precis(ZCmodel0))
a=coef(ZCmodel0)["a"]
cat("Probability of correct response: ", logistic(a))
		
# All cards and subjects treated the same
ZCmodel1 <- ulam(
			alist( Score ~ dbinom(1, p), 
			logit(p) <- a,
			a ~ dnorm(0, 10)
			), data= ZCdataClean, log_lik=TRUE, chains=4  ) 
		
print(precis(ZCmodel1))

a=coef(ZCmodel1)["a"]
cat("Probability of correct response: ", logistic(a))

if(1==1){
# Different probabilities for different participants
ZCmodel2 <- ulam(
			alist( Score ~ dbinom(1, p), 
			logit(p) <- a[Participant],
			a[Participant] ~ dnorm(0, 10)
			), data= ZCdataClean, log_lik=TRUE, chains=4  ) 
		
save(ZCmodel2, file="ZCParticipants.Rpd")	
}else{
	load("ZCParticipants.Rpd")
	}

print(precis(ZCmodel2, depth=2))

modelCoefficients <- coef(ZCmodel2)
aValues2<-c()
for(p in unique(ZCdata$Participant)){
	code<-sprintf("a[%d]", p)
	aValues2<- c(aValues2, modelCoefficients[code])	
}	
cat("Probability of correct response: ", logistic(aValues2))

if(1==1){
# Different probabilities for different participants, with shrinkage
ZCMultilevel <- ulam(
			alist( Score ~ dbinom(1, p), 
			logit(p) <- a[Participant],
			a[Participant] ~ dnorm(grand, sigma_part),
			grand ~dnorm(0, 10),
			sigma_part ~ dunif(0, 100)			
			), data= ZCdataClean, iter=1e4, warmup=2000, log_lik=TRUE, chains=4  ) 
save(ZCMultilevel, file="ZCMultilevel.Rpd")	
}else{
	load("ZCMultilevel.Rpd")
	}

		
print(precis(ZCMultilevel, depth=2))

modelCoefficients <- coef(ZCMultilevel)
aValuesMulti<-c()
for(p in unique(ZCdata $Participant)){
	code<-sprintf("a[%d]", p)
	aValuesMulti <- c(aValuesMulti, modelCoefficients[code])	
}
cat("Probability of correct response: ", logistic(aValuesMulti))

print(compare(ZCmodel0, ZCmodel1, ZCmodel2, ZCMultilevel))

# Plots to show shrinkage
participant.seq <- seq(from=1, to=22, by=1)
plot(participant.seq, logistic(aValues2), pch=1, col=col.alpha("green",1), ylab="Probability")
points(participant.seq, logistic(aValuesMulti), pch=15, col=col.alpha("red",1) )
shrinkTo <- logistic(as.numeric(coef(ZCMultilevel)["grand"]))
lines(c(0, 23), c(shrinkTo, shrinkTo), col=col.alpha("red", 1))