rm(list=ls(all=TRUE))  # clear all variables
graphics.off()  # clear all graphics
# Visual Search
# Greg Francis
# PSY 646
# 07 February 2024
# This version includes data from multiple participants with a regularizing prior on the slope and intercepts

# load the rethinking library
library(rethinking)

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

# pull out just the trials for all  participant's conjunctive target absent condition
VSdata2<-subset(VSdata, VSdata$DistractorType=="Conjunction" )
#  

# fit a linear model that predicts response time as a function of the number of distractors

VSdata2$TargetIsPresent <- ifelse(VSdata2$Target=="Present", 1, 0)  # 1 for Target present, 2 for Target Absent

# Without participant index (for homogeneous model)
cleanVSdata2<- data.frame(RT_ms=VSdata2$RT_ms, TargetIsPresent =VSdata2$TargetIsPresent, NumberDistractors=VSdata2$NumberDistractors)
			
# Twice slope, different sigma's
# Same slope and intercept for each participant
VSmodelHomogeneous1 <- map(
	alist( RT_ms ~ dnorm(mu, sigma), 
	mu <- a + (b +(1-TargetIsPresent)*b)*NumberDistractors,
	a ~ dnorm(1000, 500),
	b ~ dnorm(0, 100),
	sigma <- c1*TargetIsPresent + c2*(1-TargetIsPresent),
	c1 ~ dunif(0, 2000),
	c2 ~ dunif(0, 2000)
	), data= cleanVSdata2 ) 
cat("Finished with MAP homogeneous model")


# Same slope and intercept for each participant
VSmodelHomogeneous2 <- ulam(
	alist( RT_ms ~ dnorm(mu, sigma), 
	mu <- a + (b +(1-TargetIsPresent)*b)*NumberDistractors,
	a ~ dnorm(1000, 500),
	b ~ dnorm(0, 100),
	sigma <- c1*TargetIsPresent + c2*(1-TargetIsPresent),
	c1 ~ dunif(0, 2000),
	c2 ~ dunif(0, 2000)
	), data= cleanVSdata2, log_lik=TRUE, chains=1) 

cat("Finished with Stan homogeneous model")

print(summary(VSmodelHomogeneous1))	
print(summary(VSmodelHomogeneous2))

# Check MCMC sampling process
#traceplot(VSmodelHomogeneous2)

plot(VSmodelHomogeneous2)

dev.new()
#pairs(VSmodelHomogeneous2)

dev.new()

NumberDistractors.seq<-seq(from=1, to=65, by=1)
# Plot for MAP version
# Target present
plot(RT_ms ~ NumberDistractors, data=subset(VSdata2, cleanVSdata2$TargetIsPresent==1 ), pch=15, col=col.alpha("red",0.2))
# Target absent
points(RT_ms ~ NumberDistractors, data=subset(cleanVSdata2, cleanVSdata2$TargetIsPresent==0 ), pch=1, col=col.alpha("green",0.2))
			
abline(a=coef(VSmodelHomogeneous1)["a"], b=coef(VSmodelHomogeneous1)["b"], col=col.alpha("black",1.0))
abline(a=coef(VSmodelHomogeneous1)["a"], b=2*coef(VSmodelHomogeneous1)["b"], col=col.alpha("black",1.0))			
# Uncertainty
mu_absent1<-link(VSmodelHomogeneous1, data=data.frame(NumberDistractors=NumberDistractors.seq, TargetIsPresent =0))
mu_absent.mean1<- apply(mu_absent1$mu, 2, mean)
mu_absent.HPDI1 <-apply(mu_absent1$mu, 2, HPDI, prob=0.89)
shade(mu_absent.HPDI1, NumberDistractors.seq, col=col.alpha("green",0.4))

mu_present1<-link(VSmodelHomogeneous1, data=data.frame(NumberDistractors=NumberDistractors.seq, TargetIsPresent =1))
mu_present.mean1<- apply(mu_present1$mu, 2, mean)
mu_present.HPDI1 <-apply(mu_present1$mu, 2, HPDI, prob=0.89)
shade(mu_present.HPDI1, NumberDistractors.seq, col=col.alpha("red",0.4))
	
# Plot for Stan version		
dev.new()
# Target present
plot(RT_ms ~ NumberDistractors, data=subset(VSdata2, cleanVSdata2$TargetIsPresent==1 ), pch=15, col=col.alpha("red",0.2))
# Target absent
points(RT_ms ~ NumberDistractors, data=subset(cleanVSdata2, cleanVSdata2$TargetIsPresent==0 ), pch=1, col=col.alpha("green",0.2))
			
abline(a=coef(VSmodelHomogeneous2)["a"], b=coef(VSmodelHomogeneous2)["b"], col=col.alpha("red",1.0))
abline(a=coef(VSmodelHomogeneous2)["a"], b=2*coef(VSmodelHomogeneous2)["b"], col=col.alpha("green",1.0))			
# Uncertainty
mu_absent1<-link(VSmodelHomogeneous2, data=data.frame(NumberDistractors=NumberDistractors.seq, TargetIsPresent =0))
mu_absent.mean1<- apply(mu_absent1$mu, 2, mean)
mu_absent.HPDI1 <-apply(mu_absent1$mu, 2, HPDI, prob=0.89)
shade(mu_absent.HPDI1, NumberDistractors.seq, col=col.alpha("green",0.4))

mu_present1<-link(VSmodelHomogeneous2, data=data.frame(NumberDistractors=NumberDistractors.seq, TargetIsPresent =1))
mu_present.mean1<- apply(mu_present1$mu, 2, mean)
mu_present.HPDI1 <-apply(mu_present1$mu, 2, HPDI, prob=0.89)
shade(mu_present.HPDI1, NumberDistractors.seq, col=col.alpha("red",0.4))