rm(list=ls(all=TRUE))  # clear all variables
graphics.off()  # clear all graphics
# Visual Search
# Greg Francis
# PSY 646
# 25 September 2016
# 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)
# Dummy variable to indicate target present or absent
VSdata$ParticipantIndex <- coerce_index(VSdata$Participant)

# pull out just the trials for the first 5  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
	
# Clean up so it will work for Stan
cleanVSdata<- data.frame(ParticipantIndex =VSdata2$ParticipantIndex, RT_ms=VSdata2$RT_ms, TargetIsPresent =VSdata2$TargetIsPresent, NumberDistractors=VSdata2$NumberDistractors)

# 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
VSmodelHomogeneous <- 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  ) 

cat("Finished with homogeneous model")
print(summary(VSmodelHomogeneous))

# Multi-level intercept and slope for each participant
VSmodelMultilevel <- ulam(
	alist( RT_ms ~ dnorm(mu, sigma), 
	mu <- a[ParticipantIndex] + (b[ParticipantIndex] +(1-TargetIsPresent)*b[ParticipantIndex])*NumberDistractors,
	a[ParticipantIndex] ~ dnorm(intercept, 100),
	b[ParticipantIndex] ~ dnorm(slope, 1),
	intercept ~ dnorm(1000, 1000),
	slope ~dnorm(0, 100),
	sigma <- c1*TargetIsPresent + c2*(1-TargetIsPresent),
	c1 ~ dunif(0, 2000),
	c2 ~ dunif(0, 2000)
	), data= cleanVSdata, log_lik=TRUE, chains=4   ) 
		
cat("Finished with VSmodelMultilevel model")
print(summary(VSmodelMultilevel))


#Separate intercept and slope for each participant
VSmodelIndependent <- ulam(
	alist( RT_ms ~ dnorm(mu, sigma), 
	mu <- a[ParticipantIndex] + (b[ParticipantIndex] +(1-TargetIsPresent)*b[ParticipantIndex])*NumberDistractors,
	a[ParticipantIndex] ~ dnorm(1000, 500),
	b[ParticipantIndex] ~ dnorm(0, 100),
	sigma <- c1*TargetIsPresent + c2*(1-TargetIsPresent),
	c1 ~ dunif(0, 2000),
	c2 ~ dunif(0, 2000)
	), data= cleanVSdata, log_lik=TRUE, chains=4  ) 
cat("Finished with VSmodelIndependent model")
print(summary(VSmodelIndependent))


print(compare(VSmodelHomogeneous, VSmodelIndependent, VSmodelMultilevel))


# Multi-level intercept and slope for each participant
VSmodelMultilevel4 <- ulam(
  alist( RT_ms ~ dnorm(mu, sigma), 
         mu <- a[ParticipantIndex] + (b[ParticipantIndex] +(1-TargetIsPresent)*b[ParticipantIndex])*NumberDistractors,
         a[ParticipantIndex] ~ dnorm(intercept, sigma_part1),
         b[ParticipantIndex] ~ dnorm(slope, sigma_part2),
         intercept ~ dnorm(1000, 1000),
         slope ~dnorm(0, 100),
         sigma_part1 ~ dunif(0, 2000),
         sigma_part2 ~ dunif(0, 2000),
         sigma <- c1*TargetIsPresent + c2*(1-TargetIsPresent),
         c1 ~ dunif(0, 2000),
         c2 ~ dunif(0, 2000)
  ), data= cleanVSdata, log_lik=TRUE, chains=4) 

cat("Finished with VSmodelMultilevel model")


# pull out slopes coefficients for independent and multilevel models
mlevelb <- as.numeric(coef(VSmodelMultilevel4)[69:136])
indepb <- as.numeric(coef(VSmodelIndependent)[69:136])

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

plot(participant.seq, indepb, pch=1, col=col.alpha("green",1), ylab="slope")
points(participant.seq, mlevelb, pch=15, col=col.alpha("red",1) )
shrinkTo <- as.numeric(coef(VSmodelMultilevel4)["slope"])
lines(c(0, 70), c(shrinkTo, shrinkTo))

# pull out intercept coefficients for independent and multilevel models
mlevela <- as.numeric(coef(VSmodelMultilevel4)[1:68])
indepa <- as.numeric(coef(VSmodelIndependent)[1:68])

dev.new()
plot(participant.seq, indepa, pch=1, col=col.alpha("green",1), ylab="intercept")
points(participant.seq, mlevela, pch=15, col=col.alpha("red",1) )
shrinkToa <- as.numeric(coef(VSmodelMultilevel4)["intercept"])
lines(c(0, 70), c(shrinkToa, shrinkToa))