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


# load the rethinking library
library(rethinking)

# load full data file
Humdata<-read.csv(file="HummingRT.csv",header=TRUE,stringsAsFactors=FALSE)
# Dummy variable to indicate participant ID 
Humdata$ParticipantIndex <- coerce_index(Humdata$ID)
Humdata$TaskIndex <- coerce_index(Humdata$Task)

Humdata$IsHumming <- ifelse(Humdata$Task=="Humming", 1, 0)  # 1 when humming, 0 otherwise
Humdata$IsQuiet <- ifelse(Humdata$Task=="Quiet", 1, 0)  # 1 when humming, 0 otherwise
	
# Clean up data frame so it will work for Stan

# Without participant index (for null  model)
cleanHumdata<- data.frame(ParticipantIndex =Humdata$ParticipantIndex, ResponseTime=Humdata$ResponseTime, Task=Humdata$Task)


# Differences across participants, no difference across conditions
HumNull <- ulam(
	alist( ResponseTime ~ dnorm(mu, sigma), 
	mu <- a[ParticipantIndex],
	a[ParticipantIndex] ~ dnorm(500, 200),
	sigma ~ dunif(0, 2000)
	), data= cleanHumdata, iter=5000, warmup=2000, log_lik=TRUE, chains=4  ) 

cat("Finished with null model")
print(summary(HumNull))


cleanHumdata2<- data.frame(ParticipantIndex =Humdata$ParticipantIndex, ResponseTime=Humdata$ResponseTime, TaskIndex=Humdata$TaskIndex, IsQuiet = Humdata$IsQuiet)

#  Differences across participants and difference across conditions
#  Use Quiet as the "baseline" condition, from which other conditions might deviate
HumFull <- ulam(
	alist( ResponseTime ~ dnorm(mu, sigma), 
	mu <- a[ParticipantIndex] + b[TaskIndex]*(1-IsQuiet),
	a[ParticipantIndex] ~ dnorm(500, 200),
	b[TaskIndex] ~ dnorm(0,20),	
	sigma ~ dunif(0, 2000)
	), data= cleanHumdata2, iter=10000, warmup=5000, log_lik=TRUE, chains=4   ) 
		
cat("Finished with full model")
print(summary(HumFull))

cleanHumdata3 <- data.frame(ParticipantIndex =Humdata$ParticipantIndex, ResponseTime=Humdata$ResponseTime, TaskIndex=Humdata$TaskIndex, IsHumming = Humdata$IsHumming)

#  Differences across participants, predicted effects across conditions
HumPredicted <- ulam(
	alist( ResponseTime ~ dnorm(mu, sigma), 
	mu <- a[ParticipantIndex] -b*IsHumming,
	a[ParticipantIndex] ~ dnorm(500, 200),
	b ~ dunif(0, 30),
	sigma ~ dunif(0, 2000)
	), data= cleanHumdata3, iter=10000, warmup=5000, log_lik=TRUE, chains=4  ) 
cat("Finished with predicted model")
print(summary(HumPredicted))


print(compare(HumNull, HumFull, HumPredicted))


# Multi-level model that applies shrinkage to participant mean
HumMultiLevel <- ulam(
  alist( ResponseTime ~ dnorm(mu, sigma), 
         mu <- a[ParticipantIndex] + b[TaskIndex]*(1-IsQuiet),
		 a[ParticipantIndex] ~ dnorm(parmean, parsd),
         parmean ~ dnorm(500, 200),
         parsd ~ dunif(0, 2000),
		 b[TaskIndex] ~ dnorm(0, 20),	
         sigma ~ dunif(0, 2000)
  ), data= cleanHumdata2 , iter=40000, warmup=20000, log_lik=TRUE, chains=4 ) 

cat("Finished with HumMultiLevel model")
print(summary(HumMultiLevel))

print(compare(HumNull, HumFull, HumPredicted, HumMultiLevel))


# pull out coefficients for independent and multilevel models
mlevela <- as.numeric(coef(HumMultiLevel)[1:104])
indepa <- as.numeric(coef(HumFull)[1:104])

participant.seq <- seq(from=1, to=104, by=1)

plot(participant.seq, indepa, pch=1, col=col.alpha("green",1), ylab="Participant mean", ylim=c(0,500))
points(participant.seq, mlevela, pch=15, col=col.alpha("red",1) )
shrinkTo <- as.numeric(coef(HumMultiLevel)["parmean"])
lines(c(0, 105), c(shrinkTo, shrinkTo))