rm(list=ls(all=TRUE))  # clear all variables
graphics.off()  # clear all graphics
# Map Search (Francis, Bias & Shive, 2010)
# Greg Francis
# PSY 626
# 22 February 2024

# Using Stan

# load the rethinking library
library(rethinking)

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

MSdata1<- data.frame(RT_ms= MSdata$RT_ms, Proximity = MSdata$Proximity, Size = MSdata$Size, Contrast = MSdata$Contrast)

# fit a linear model that predicts response time as a function of the other terms (basically same as done in original paper)
MSLinearModel <- ulam(
	alist( RT_ms ~ dnorm(mu, sigma), 
	mu <- a0 + a1*Proximity + a2*Size + a3*Color + a4*Contrast,
	a0 ~ dnorm(1000, 1000),
	a1 ~ dnorm(0, 20),
	a2 ~ dnorm(0, 50),
	a3 ~ dnorm(0, 1),
	a4 ~ dnorm(0, 500),
	sigma  ~ dunif(0, 2000)
	), data= MSdata, log_lik=TRUE, chains=4 ) 

cat("Finished with linear model\n")
print(summary(MSLinearModel))

# Linear without color
MSLinearNoColorModel <- ulam(
	alist( RT_ms ~ dnorm(mu, sigma), 
	mu <- a0 + a1*Proximity + a2*Size + a4*Contrast,
	a0 ~ dnorm(1000, 1000),
	a1 ~ dnorm(0, 20),
	a2 ~ dnorm(0, 50),
	a4 ~ dnorm(0, 500),
	sigma  ~ dunif(0, 2000)
	), data= MSdata1, log_lik=TRUE, chains=4 ) 

cat("\n\nFinished with linear w/o color model\n")
print(summary(MSLinearNoColorModel))

# interaction model with Color and Contrast
MSColorContrastInteractionModel <- ulam(
	alist( RT_ms ~ dnorm(mu, sigma), 
	mu <- a0 + a1*Proximity + a2*Size + a3*Color + a4*Contrast + b3*Color*Contrast,
	a0 ~ dnorm(1000, 1000),
	a1 ~ dnorm(0, 20),
	a2 ~ dnorm(0, 50),
	a3 ~ dnorm(0, 50), 
	b3 ~ dnorm(0, 1),
	a4 ~ dnorm(0, 500),
	sigma  ~ dunif(0, 2000)
	), data= MSdata, log_lik=TRUE, chains=4  ) 

cat("\n\nFinished with color/contrast interaction model\n")
print(summary(MSColorContrastInteractionModel))

# interaction model with Size and Contrast
MSSizeContrastInteractionModel <- ulam(
	alist( RT_ms ~ dnorm(mu, sigma), 
	mu <- a0+ a1*Proximity + a2*Size + a3*Color + a4*Contrast + b2*Size*Contrast,
	a0 ~ dnorm(1000, 1000),
	a1 ~ dnorm(0, 20),
	a2 ~ dnorm(0, 50),
	a3 ~ dnorm(0, 1),
	a4 ~ dnorm(0, 500),
	b2 ~ dnorm(0, 500),
	sigma  ~ dunif(0, 2000)
	), data= MSdata, log_lik=TRUE, chains=4 ) 

cat("\n\nFinished with size/contrast interaction model\n")
print(summary(MSSizeContrastInteractionModel))

# interaction model with Contrast interacting with everything
MSAllContrastInteractionModel <- ulam(
	alist( RT_ms ~ dnorm(mu, sigma), 
	mu <- a0+ a1*Proximity + a2*Size + a3*Color + a4*Contrast +b1*Proximity*Contrast + b2*Size*Contrast + b3*Size*Color,
	a0 ~ dnorm(1000, 1000),
	a1 ~ dnorm(0, 20),
	a2 ~ dnorm(0, 50),
	a3 ~ dnorm(0, 50),
	a4 ~ dnorm(0, 500),	
	b1 ~ dnorm(0, 500),
	b2 ~ dnorm(0, 500),
	b3 ~ dnorm(0, 500),
	sigma  ~ dunif(0, 2000)
	), data= MSdata, log_lik=TRUE, chains=4 ) 

cat("\n\nFinished with all/contrast interaction model\n")
print(summary(MSAllContrastInteractionModel))

# Compare models
compareModels <- compare(MSLinearModel, MSLinearNoColorModel, MSColorContrastInteractionModel, MSSizeContrastInteractionModel, MSAllContrastInteractionModel)

print(compare(compareModels) )


# SizeContrastInteraction model seems to be the best. Use it to predict RT for a word cloud

mu<-link(MSSizeContrastInteractionModel, data=data.frame(Color=c(40280.93, 32924.34, 48904.35), Proximity=c(94.2, 94.2, 94.2), Size=c(-0.9507, -0.9507, -0.9507), Contrast=c(0.3406859, 0.6552078, 1)))
mu.mean <- apply(mu, 2, mean)
mu.HPDI <-apply(mu, 2, HPDI, prob=0.89)


# Compute expected search time weighted by model probability
# Extract weights and model order from compare result
modelWeights <- compareModels[[6]]
orderModelNames <- row.names(compareModels)

sumST=0

# Linear model
mu<-link(MSLinearModel, data=data.frame(Color=c(40280.93, 32924.34, 48904.35), Proximity=c(94.2, 94.2, 94.2), Size=c(-0.9507, -0.9507, -0.9507), Contrast=c(0.3406859, 0.6552078, 1)))
mu.mean <- apply(mu, 2, mean)
# Find which weight to apply
thisWeight=0
for(i in c(1:length(orderModelNames))){
	if(orderModelNames[i]=="MSLinearModel"){thisWeight=modelWeights[i]}
}	
sumST = sumST + thisWeight*mu.mean

# MSLinearNoColorModel
mu<-link(MSLinearNoColorModel, data=data.frame(Color=c(40280.93, 32924.34, 48904.35), Proximity=c(94.2, 94.2, 94.2), Size=c(-0.9507, -0.9507, -0.9507), Contrast=c(0.3406859, 0.6552078, 1)))
mu.mean <- apply(mu, 2, mean)
# Find which weight to apply
thisWeight=0
for(i in c(1:length(orderModelNames))){
	if(orderModelNames[i]=="MSLinearNoColorModel"){thisWeight=modelWeights[i]}
}	
sumST = sumST + thisWeight*mu.mean

# MSColorContrastInteractionModel
mu<-link(MSColorContrastInteractionModel, data=data.frame(Color=c(40280.93, 32924.34, 48904.35), Proximity=c(94.2, 94.2, 94.2), Size=c(-0.9507, -0.9507, -0.9507), Contrast=c(0.3406859, 0.6552078, 1)))
mu.mean <- apply(mu, 2, mean)
# Find which weight to apply
thisWeight=0
for(i in c(1:length(orderModelNames))){
	if(orderModelNames[i]=="MSColorContrastInteractionModel"){thisWeight=modelWeights[i]}
}	
sumST = sumST + thisWeight*mu.mean

# MSSizeContrastInteractionModel
mu<-link(MSSizeContrastInteractionModel, data=data.frame(Color=c(40280.93, 32924.34, 48904.35), Proximity=c(94.2, 94.2, 94.2), Size=c(-0.9507, -0.9507, -0.9507), Contrast=c(0.3406859, 0.6552078, 1)))
mu.mean <- apply(mu, 2, mean)
# Find which weight to apply
thisWeight=0
for(i in c(1:length(orderModelNames))){
	if(orderModelNames[i]=="MSSizeContrastInteractionModel"){thisWeight=modelWeights[i]}
}	
sumST = sumST + thisWeight*mu.mean

# MSAllContrastInteractionModel
mu<-link(MSAllContrastInteractionModel, data=data.frame(Color=c(40280.93, 32924.34, 48904.35), Proximity=c(94.2, 94.2, 94.2), Size=c(-0.9507, -0.9507, -0.9507), Contrast=c(0.3406859, 0.6552078, 1)))
mu.mean <- apply(mu, 2, mean)
# Find which weight to apply
thisWeight=0
for(i in c(1:length(orderModelNames))){
	if(orderModelNames[i]=="MSAllContrastInteractionModel"){thisWeight=modelWeights[i]}
}	
sumST = sumST + thisWeight*mu.mean

cat("Model weighted predicted search times are:", sumST, "\n")