rm(list=ls(all=TRUE))  # clear all variables
graphics.off()  # clear all graphics
# ADHD data set
# Greg Francis
# PSY 626
# 01 February 2024

# load data file
ADdata<-read.csv(file="ADHD.csv", header=TRUE,stringsAsFactors=FALSE)

# load the rethinking library
library(rethinking)

# Dummy variable to indicate the 4 doses
ADdata$ConditionIndex <-0*c(1:4)
ADdata$ConditionIndex[ADdata$Dose =="D0"] =1
ADdata$ConditionIndex[ADdata$Dose =="D15"] =2
ADdata$ConditionIndex[ADdata$Dose =="D30"] =3
ADdata$ConditionIndex[ADdata$Dose =="D60"] =4


# Null model: one mean for all conditions
ADmodel0 <- map(
			alist( Delay ~ dnorm(mu, sigma), 
			mu <- a[Participant],
			a[Participant] ~ dnorm(50, 50),
			sigma ~ dunif(0, 100)
			), data= ADdata , control=list(maxit=10000) ) 
cat("Finished ADmodel0\n")
			
			
# Alternative model 1: different mean for each condition	
ADmodel1 <- map(
			alist( Delay ~ dnorm(mu, sigma), 
			mu <- a[Participant] +b[ConditionIndex],
			a[Participant] ~ dnorm(50, 50),
			b[ConditionIndex] ~ dnorm(0, 100),
			sigma ~ dunif(0, 100)
			), data= ADdata , control=list(maxit=10000) ) 

cat("Finished ADmodel2\n")
	
plot(coeftab(ADmodel0, ADmodel1 ))
print(compare(ADmodel0, ADmodel1))

# Characterize the posterior distributions
post<-extract.samples(ADmodel1, n= 10000)

# Compute posterior of mean (across participants) for each dosage

D0.Means <- apply(post$a + post$b[,1], 1, mean)
D15.Means <- apply(post$a + post$b[,2], 1, mean)
D30.Means <- apply(post$a + post$b[,3], 1, mean)
D60.Means <- apply(post$a + post$b[,4], 1, mean)


dev.new()
dens(D0.Means, col=rangi2, lwd=2, xlab="posterior for D0")
dev.new()
dens(D15.Means, col=rangi2, lwd=2, xlab="posterior for D15")
dev.new()
dens(D30.Means, col=rangi2, lwd=2, xlab="posterior for D30")
dev.new()
dens(D60.Means, col=rangi2, lwd=2, xlab="posterior for D60")

# compute some interesting probabilities (kind of like contrasts)
cat("\nProbability that mean for D15 is greater than mean for D0 is: ", length(D15.Means[D15.Means>D0.Means])/length(D15.Means), "\n")
cat("Probability that mean for D30 is greater than mean for D0 is: ", length(D30.Means[D30.Means>D0.Means])/length(D30.Means), "\n")

cat("Probability that mean for D60 is greater than mean for D0 is: ", length(D60.Means[D60.Means>D0.Means])/length(D60.Means), "\n")

ComboLow = D0.Means + D15.Means
ComboHigh = D30.Means + D60.Means
cat("Probability that mean for D0+D15 is less than mean for D30+D60 is: ", length(ComboLow[ComboLow<ComboHigh])/length(ComboLow), "\n")


#--------Exploring priors

# Tighter prior for b terms
ADmodel2 <- map(
			alist( Delay ~ dnorm(mu, sigma), 
			mu <- a[Participant] +b[ConditionIndex],
			a[Participant] ~ dnorm(50, 50),
			b[ConditionIndex] ~ dnorm(0, 10),
			sigma ~ dunif(0, 100)
			), data= ADdata , control=list(maxit=10000) ) 


# Even tighter prior for b terms
ADmodel3 <- map(
			alist( Delay ~ dnorm(mu, sigma), 
			mu <- a[Participant] +b[ConditionIndex],
			a[Participant] ~ dnorm(50, 50),
			b[ConditionIndex] ~ dnorm(0, 1),
			sigma ~ dunif(0, 100)
			), data= ADdata , control=list(maxit=10000) )