rm(list=ls(all=TRUE))  # clear all variables
graphics.off()  # clear all graphics
# Sternberg Search
# Greg Francis
# PSY 626
# March 27, 2024


# load the rethinking library
library(rethinking)

# load full data file
SSdata<-read.csv(file="SternbergSearch.csv",header=TRUE,stringsAsFactors=FALSE)
SSdata$TargetPresent <- ifelse(SSdata$Condition =="Present", 1, 0)  


# Stan model for null (just intercept), with shrinkage
SSdataNull <- data.frame(RT= SSdata$RT, Participant=SSdata$Participant)		
SSmodelNull <- ulam(
			alist( RT ~ dnorm(mu, sigma), 
			mu <- a[Participant],
			a[Participant] ~ dnorm(grand_mu, grand_s),
			grand_mu ~ dnorm(1000, 1000),
			grand_s ~ dunif(0, 2000),
			sigma ~ dunif(0, 1000)
			), data= SSdataNull, log_lik=TRUE, iter=10000, warmup=4000, chains=4, cores=4  ) 
				
cat("Finished SSmodelNull\n")

# Stan model for main effect of MemorySetSize, with shrinkage
SSdataMSS <- data.frame(RT= SSdata$RT, MemorySetSize=SSdata$MemorySetSize, Participant=SSdata$Participant)		
SSmodelMSS <- ulam(
			alist( RT ~ dnorm(mu, sigma), 
			mu <- a[Participant] + b[Participant]*MemorySetSize,
			a[Participant] ~ dnorm(grand_mua, grand_sa),
			b[Participant] ~ dnorm(grand_mub, grand_sb),
			grand_mua ~ dnorm(1000, 1000),
			grand_sa ~ dunif(0, 2000),
			grand_mub ~ dnorm(0, 100),
			grand_sb ~ dunif(0, 200),			
			sigma ~ dunif(0, 1000)
			), data= SSdataMSS, log_lik=TRUE, iter=10000, warmup=4000, chains=4, cores=4  ) 
				
cat("Finished SSmodelMSS\n")



# Stan model for main effect of Target condition, with shrinkage
SSdataCondition <- data.frame(RT= SSdata$RT, Condition=SSdata$TargetPresent, Participant=SSdata$Participant)		
SSmodelCondition <- ulam(
			alist( RT ~ dnorm(mu, sigma), 
			mu <- a[Participant] + b[Participant]*Condition,
			a[Participant] ~ dnorm(grand_mua, grand_sa),
			b[Participant] ~ dnorm(grand_mub, grand_sb),
			grand_mua ~ dnorm(1000, 1000),
			grand_sa ~ dunif(0, 2000),
			grand_mub ~ dnorm(0, 100),
			grand_sb ~ dunif(0, 200),
			sigma ~ dunif(0, 1000)
			), data= SSdataCondition, log_lik=TRUE, iter=10000, warmup=4000, chains=4, cores=4  ) 
				
cat("Finished SSmodelCondition\n")


# Stan model for additive effects 
SSdataAdditive <- data.frame(RT= SSdata$RT, Condition=SSdata$TargetPresent, MemorySetSize=SSdata$MemorySetSize, Participant=SSdata$Participant)		
SSmodelAdditive <- ulam(
			alist( RT ~ dnorm(mu, sigma), 
			mu <- a[Participant] + b[Participant]*MemorySetSize + c[Participant]*Condition,
			a[Participant] ~ dnorm(grand_mua, grand_sa),
			b[Participant] ~ dnorm(grand_mub, grand_sb),
			c[Participant] ~ dnorm(grand_muc, grand_sc),
			grand_mua ~ dnorm(1000, 1000),
			grand_sa ~ dunif(0, 2000),
			grand_mub ~ dnorm(0, 100),
			grand_sb ~ dunif(0, 200),
			grand_muc ~ dnorm(0, 100),
			grand_sc ~ dunif(0, 200),
			sigma ~ dunif(0, 1000)
			), data= SSdataAdditive, log_lik=TRUE, iter=10000, warmup=4000, chains=4, cores=4  ) 
				
cat("Finished SSmodelAdditive\n")

# Stan model for additive effects and interaction
SSdataInteraction <- data.frame(RT= SSdata$RT, Condition=SSdata$TargetPresent, MemorySetSize=SSdata$MemorySetSize, Participant=SSdata$Participant)		
SSmodelInteraction <- ulam(
			alist( RT ~ dnorm(mu, sigma), 
			mu <- a[Participant] + b[Participant]*MemorySetSize + c[Participant]*Condition+ d[Participant]*MemorySetSize*Condition,
			a[Participant] ~ dnorm(grand_mua, grand_sa),
			b[Participant] ~ dnorm(grand_mub, grand_sb),
			c[Participant] ~ dnorm(grand_muc, grand_sc),
			d[Participant] ~ dnorm(grand_mud, grand_sd),
			grand_mua ~ dnorm(1000, 1000),
			grand_sa ~ dunif(0, 2000),
			grand_mub ~ dnorm(0, 100),
			grand_sb ~ dunif(0, 200),
			grand_muc ~ dnorm(0, 100),
			grand_sc ~ dunif(0, 200),
			grand_mud ~ dnorm(0, 100),
			grand_sd ~ dunif(0, 200),
			sigma ~ dunif(0, 1000)
			), data= SSdataInteraction, log_lik=TRUE, iter=10000, warmup=4000, chains=4, cores=4  ) 
				
cat("Finished SSmodelInteraction\n")


# Stan model for interaction only (no additive effects)
SSdataInteraction <- data.frame(RT= SSdata$RT, Condition=SSdata$TargetPresent, MemorySetSize=SSdata$MemorySetSize, Participant=SSdata$Participant)		
SSmodelInteractionOnly <- ulam(
			alist( RT ~ dnorm(mu, sigma), 
			mu <- a[Participant] + d[Participant]*MemorySetSize*Condition,
			a[Participant] ~ dnorm(grand_mua, grand_sa),
			d[Participant] ~ dnorm(grand_mud, grand_sd),
			grand_mua ~ dnorm(1000, 1000),
			grand_sa ~ dunif(0, 2000),
			grand_mud ~ dnorm(0, 100),
			grand_sd ~ dunif(0, 200),
			sigma ~ dunif(0, 1000)
			), data= SSdataInteraction, log_lik=TRUE, iter=10000, warmup=4000, chains=4, cores=4  ) 
	
cat("Finished SSmodelInteractionOnly\n")
				
	
# Stan model for Memory Set size and interaction
SSdataInteraction <- data.frame(RT= SSdata$RT, Condition=SSdata$TargetPresent, MemorySetSize=SSdata$MemorySetSize, Participant=SSdata$Participant)		
SSmodelInteractionMMS <- ulam(
			alist( RT ~ dnorm(mu, sigma), 
			mu <- a[Participant] + b[Participant]*MemorySetSize + d[Participant]*MemorySetSize*Condition,
			a[Participant] ~ dnorm(grand_mua, grand_sa),
			b[Participant] ~ dnorm(grand_mub, grand_sb),
			d[Participant] ~ dnorm(grand_mud, grand_sd),
			grand_mua ~ dnorm(1000, 1000),
			grand_sa ~ dunif(0, 2000),
			grand_mub ~ dnorm(0, 100),
			grand_sb ~ dunif(0, 200),
			grand_mud ~ dnorm(0, 100),
			grand_sd ~ dunif(0, 200),
			sigma ~ dunif(0, 1000)
			), data= SSdataInteraction, log_lik=TRUE, iter=10000, warmup=4000, chains=4, cores=4  ) 
					
cat("Finished SSmodelInteractionMMS\n")

# Stan model for target condition and interaction
SSdataInteraction <- data.frame(RT= SSdata$RT, Condition=SSdata$TargetPresent, MemorySetSize=SSdata$MemorySetSize, Participant=SSdata$Participant)		
SSmodelInteractionCondition <- ulam(
			alist( RT ~ dnorm(mu, sigma), 
			mu <- a[Participant] + c[Participant]*Condition+ d[Participant]*MemorySetSize*Condition,
			a[Participant] ~ dnorm(grand_mua, grand_sa),
			c[Participant] ~ dnorm(grand_muc, grand_sc),
			d[Participant] ~ dnorm(grand_mud, grand_sd),
			grand_mua ~ dnorm(1000, 1000),
			grand_sa ~ dunif(0, 2000),
			grand_muc ~ dnorm(0, 100),
			grand_sc ~ dunif(0, 200),
			grand_mud ~ dnorm(0, 100),
			grand_sd ~ dunif(0, 200),
			sigma ~ dunif(0, 1000)
			), data= SSdataInteraction, log_lik=TRUE, iter=10000, warmup=4000, chains=4, cores=4  ) 
	
cat("Finished SSmodelInteractionCondition\n")
				
		
print(compare(SSmodelInteractionCondition, SSmodelInteractionMMS, SSmodelInteractionOnly, SSmodelInteraction, SSmodelAdditive , SSmodelCondition, SSmodelMSS, SSmodelNull))