rm(list=ls(all=TRUE))  # clear all variables
graphics.off()  # clear all graphics
# ADHD data set
# Greg Francis
# PSY 626
# 12 September 2025

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

# load the rethinking library
library(rethinking)

# Create indices for the different conditions (other than D0)
ADdata$IndexD15 <-0* ADdata$Delay  # makes an array of 0's of the correct size
ADdata$IndexD15[ADdata$Dose =="D15"] =1
ADdata$IndexD30 <-0* ADdata$Delay
ADdata$IndexD30[ADdata$Dose =="D30"] =1
ADdata$IndexD60 <-0* ADdata$Delay
ADdata$IndexD60[ADdata$Dose =="D60"] =1


# Null model: one mean for all conditions
ADmodel0 <- quap(
			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: intercept is mean for D0, slopes are deviations from D0 for other conditions



ADmodel1 <- quap(
			alist( Delay ~ dnorm(mu, sigma), 
			mu <- a[Participant] +b1*IndexD15 + b2*IndexD30 + b3*IndexD60,
			a[Participant] ~ dnorm(50, 50),
			b1 ~ dnorm(0, 100),
			b2 ~ dnorm(0, 100),
			b3 ~ 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, 1, mean)
D15.Means <- apply(post$a + post$b1, 1, mean)
D30.Means <- apply(post$a + post$b2, 1, mean)
D60.Means <- apply(post$a + post$b3, 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 <- quap(
			alist( Delay ~ dnorm(mu, sigma), 
			mu <- a[Participant] +b1*IndexD15 + b2*IndexD30 + b3*IndexD60,
			a[Participant] ~ dnorm(50, 50),
			b1 ~ dnorm(0, 10),
			b2 ~ dnorm(0, 10),
			b3 ~ 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] +b1*IndexD15 + b2*IndexD30 + b3*IndexD60,
			a[Participant] ~ dnorm(50, 50),
			b1 ~ dnorm(0, 1),
			b2 ~ dnorm(0, 1),
			b3 ~ dnorm(0, 1),
			sigma ~ dunif(0, 100)
			), data= ADdata , control=list(maxit=10000) ) 


print(compare(ADmodel0, ADmodel1, ADmodel2, ADmodel3, n=1e5))