rm(list=ls(all=TRUE))  # clear all variables
graphics.off()  # clear all graphics
# Visual Search
# Greg Francis
# PSY 646
# 20 September 2016

# 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$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")

# Dummy variable to indicate target present or absent
VSdata2$TargetIsPresent <- ifelse(VSdata2$Target=="Present", 1, 0)

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

plot(RT_ms ~ NumberDistractors, data=subset(VSdata2, VSdata2$Target=="Absent" ), pch=1)
points(RT_ms ~ NumberDistractors, data=subset(VSdata2, VSdata2$Target=="Present" ), pch=15)
abline(a=coef(VSmodel)["a"], b=coef(VSmodel)["b"], col=col.alpha("red",1.0))
abline(a=coef(VSmodel)["a"], b=coef(VSmodel)["b2"], col=col.alpha("green",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("red",0.3), lty=5)
	abline(a=post$a[i], b=post$b2[i], col=col.alpha("green",0.3), lty=5)
}

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

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

# Plot HPDI for TargetAbsent
dev.new()
plot(RT_ms ~ NumberDistractors, data=subset(VSdata2, VSdata2$Target=="Absent" ), pch=1)

# 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_absent<-link(VSmodel, data=data.frame(NumberDistractors=NumberDistractors.seq, TargetIsPresent=0))
mu_absent.mean <- apply(mu_absent, 2, mean)
mu_absent.HPDI <-apply(mu_absent, 2, HPDI, prob=0.89)
# Plot the MAP line (same as abline done previously from the linear regression coefficients)
lines(NumberDistractors.seq, mu_absent.mean, )
shade(mu_absent.HPDI, NumberDistractors.seq,  col=col.alpha("green",0.3))

# Plot HPDI for TargetPresent
points(RT_ms ~ NumberDistractors, data=subset(VSdata2, VSdata2$Target=="Present" ), pch=15)
# use link to compute mu for each sample from posterior and for each value in NumberDistractors.seq
mu_present<-link(VSmodel, data=data.frame(NumberDistractors=NumberDistractors.seq, TargetIsPresent=1))
mu_present.mean <- apply(mu_present, 2, mean)
mu_present.HPDI <-apply(mu_present, 2, HPDI, prob=0.89)
# Plot the MAP line (same as abline done previously from the linear regression coefficients)
lines(NumberDistractors.seq, mu_present.mean)
shade(mu_present.HPDI, NumberDistractors.seq, col=col.alpha("red",0.3))

# Prediction interval for RT_ms raw scores

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

# 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=subset(VSdata2, VSdata2$Target=="Absent" ), pch=1)
lines(NumberDistractors.seq, mu_absent.mean)  # MAP line for means
shade(mu_absent.HPDI, NumberDistractors.seq)   # shaded HPDI for estimates of means
shade(RT_ms.PI, NumberDistractors.seq,  col=col.alpha("green",0.3))  # shaped prediction interval for simulated RT_ms values

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

# 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
points(RT_ms ~ NumberDistractors, data=subset(VSdata2, VSdata2$Target=="Present" ), pch=15)
lines(NumberDistractors.seq, mu_present.mean)  # MAP line for means
shade(mu_present.HPDI, NumberDistractors.seq)   # shaded HPDI for estimates of means
shade(RT_ms.PI, NumberDistractors.seq,  col=col.alpha("red",0.3))  # shaped prediction interval for simulated RT_ms values