rm(list=ls(all=TRUE))  # clear all variables
graphics.off()  # clear all graphics
# Visual Search
# Greg Francis
# PSY 646
# 15 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$DistractorType=="Conjunction")
#  
# load the rethinking library
library(rethinking)

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

# Dummy variable to indicate target present or absent
VSdata2$TargetIsPresent <- ifelse(VSdata2$Target=="Present", 1, 0)  # 1 for Target present, 0 for Target Absent
			
# Different slopes, same sigma
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, control=list(maxit=10000)  ) 
		
		
dev.new()
# Plot model
# Target absent
plot(RT_ms ~ NumberDistractors, data=subset(VSdata2, VSdata2$Target=="Absent" ), pch=1, main="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_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, col=col.alpha("green",1))
lines(NumberDistractors.seq, mu_absent.HPDI[1,],   col=col.alpha("green",1), lty=2)
lines(NumberDistractors.seq, mu_absent.HPDI[2,],   col=col.alpha("green",1), lty=2)
			
shade(mu_absent.HPDI, NumberDistractors.seq, col=col.alpha("green",0.3))
			

# Target present
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, col=col.alpha("red",1))
lines(NumberDistractors.seq, mu_present.HPDI[1,],   col=col.alpha("red",1), lty=2)
lines(NumberDistractors.seq, mu_present.HPDI[2,],   col=col.alpha("red",1), lty=2)

shade(mu_present.HPDI, NumberDistractors.seq, col=col.alpha("red",0.3))

#################################		
# Twice slope, same sigma		
VSmodel2 <- map(
			alist( RT_ms ~ dnorm(mu, sigma), 
			mu <- a + b*(2-TargetIsPresent)*NumberDistractors,
			a ~ dnorm(1000, 500),
			b~ dnorm(0, 100),
			sigma ~ dunif(0, 2000)
			), data= VSdata2, control=list(maxit=10000)   ) 


dev.new()
# Target absent
plot(RT_ms ~ NumberDistractors, data=subset(VSdata2, VSdata2$Target=="Absent" ), pch=1, main="VSmodel2")
# 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(VSmodel2, 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, col=col.alpha("green",1))
lines(NumberDistractors.seq, mu_absent.HPDI[1,],   col=col.alpha("green",1), lty=2)
lines(NumberDistractors.seq, mu_absent.HPDI[2,],   col=col.alpha("green",1), lty=2)
			
shade(mu_absent.HPDI, NumberDistractors.seq, col=col.alpha("green",0.3))
			

# Target present
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(VSmodel2, 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, col=col.alpha("red",1))
lines(NumberDistractors.seq, mu_present.HPDI[1,],   col=col.alpha("red",1), lty=2)
lines(NumberDistractors.seq, mu_present.HPDI[2,],   col=col.alpha("red",1), lty=2)

shade(mu_present.HPDI, NumberDistractors.seq, col=col.alpha("red",0.3))

###################################
# Twice slope, different sigmas			
VSmodel3 <- 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= VSdata2, control=list(maxit=10000)  ) 
	


dev.new()
# Target absent
plot(RT_ms ~ NumberDistractors, data=subset(VSdata2, VSdata2$Target=="Absent" ), pch=1, main="VSmodel3")
# 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(VSmodel3, data=data.frame(NumberDistractors=NumberDistractors.seq, TargetIsPresent=0))
mu_absent.mean <- apply(mu_absent$mu, 2, mean)
mu_absent.HPDI <-apply(mu_absent$mu, 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, col=col.alpha("green",1))
lines(NumberDistractors.seq, mu_absent.HPDI[1,],   col=col.alpha("green",1), lty=2)
lines(NumberDistractors.seq, mu_absent.HPDI[2,],   col=col.alpha("green",1), lty=2)
			
shade(mu_absent.HPDI, NumberDistractors.seq, col=col.alpha("green",0.3))
			


# Target present
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(VSmodel3, data=data.frame(NumberDistractors=NumberDistractors.seq, TargetIsPresent=1))
mu_present.mean <- apply(mu_present$mu, 2, mean)
mu_present.HPDI <-apply(mu_present$mu, 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, col=col.alpha("red",1))
lines(NumberDistractors.seq, mu_present.HPDI[1,],   col=col.alpha("red",1), lty=2)
lines(NumberDistractors.seq, mu_present.HPDI[2,],   col=col.alpha("red",1), lty=2)

shade(mu_present.HPDI, NumberDistractors.seq, col=col.alpha("red",0.3))

#############
# Model comparison

print(compare(VSmodel, VSmodel2, VSmodel3))