rm(list=ls(all=TRUE))  # clear all variables
graphics.off()  # clear all graphics
# Visual Search
# Greg Francis
# PSY 626
# 14 September 2020

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

# pull out just the trials for the first participant's conjunctive target absent condition
VSdata2<-subset(VSdata, VSdata$Participant=="Francis200S16-2" & VSdata$Target=="Absent" & VSdata$DistractorType=="Conjunction")
#  
# load the rethinking library
library(rethinking)

# fit a linear model that predicts response time as a function of the number of distractors
print("LINEAR MODEL")

# Center NumberDistractors
VSdata2$NumberDistractors.c <- VSdata2$NumberDistractors - mean(VSdata2$NumberDistractors)

VSmodel <- map(
			alist( RT_ms ~ dnorm(mu, sigma), 
			mu <- a + b*NumberDistractors,
			a ~ dnorm(1000, 500),
			b ~ dnorm(0, 100),
			sigma ~ dunif(0, 500)
			), data=VSdata2 ) 
			
summaryTable <- summary(VSmodel)
print(VSmodel)		
print(precis(VSmodel, corr=TRUE))	

plot(RT_ms ~ NumberDistractors, data=VSdata2)
abline(a=coef(VSmodel)["a"], b=coef(VSmodel)["b"], col=col.alpha("red",1.0))

numVariableLines=10000
numVariableLinesToPlot=20
post<-extract.samples(VSmodel, n= numVariableLines)
for(i in 1: numVariableLinesToPlot){
	abline(a=post$a[i], b=post$b[i], col=col.alpha("black",0.3))
}

mu_at_35 <- post$a +post$b *35

dev.new()
dens(mu_at_35, col=rangi2, lwd=2, xlab="mu|NumDistract=35")

# Plot HPDI 
dev.new()
plot(RT_ms ~ NumberDistractors, data=VSdata2)

mu<-link(VSmodel)

# Define a sequence of NumberDistractors to compute predictions
NumberDistractors.seq<-seq(from=1, to=65, by=1)

# use link to compute mu for each sample from posterior and for each value in NumberDistractors.seq
mu<-link(VSmodel, data=data.frame(NumberDistractors=NumberDistractors.seq))

mu.mean <- apply(mu, 2, mean)
mu.HPDI <-apply(mu, 2, HPDI, prob=0.89)

# Plot the MAP line (same as abline done previously from the linear regression coefficients)
lines(NumberDistractors.seq, mu.mean)
shade(mu.HPDI, NumberDistractors.seq)

# Prediction interval for RT_ms raw scores

# generate many sample RT_ms scores for NumberDistractors.seq using the model
sim.RT_ms <- sim(VSmodel, data=list(NumberDistractors =NumberDistractors.seq))

# Idenitfy limits of middle 89% of samples values for each NumberDistractors (PI is a function from the rethinking library)
RT_ms.PI <- apply(sim.RT_ms, 2, PI, prob=0.89)

# Plot
dev.new()
plot(RT_ms ~ NumberDistractors, data=VSdata2)
lines(NumberDistractors.seq, mu.mean)  # MAP line for means
shade(mu.HPDI, NumberDistractors.seq)   # shaded HPDI for estimates of means
shade(RT_ms.PI, NumberDistractors.seq)  # shaped prediction interval for simulated RT_ms values