library(rjags)

## Loading required package: coda

## linking to JAGS 3.3.0

## module basemod loaded

## module bugs loaded

# Reaction time for subject 48 in ms
rt_subject_48 <- c(288, 168, 195, 198, 147, 176, 151, 230, 180, 190, 199, 150, 
    176, 197, 216, 232, 249, 280, 210, 200, 322, 303, 317, 243, 266, 
    296, 206, 175, 190, 167)

# I have not though about the priors at all, I'm just using flat priors for now 
# for convenience.


###
### A standard normal distributed model...
###

norm_model_str <- "
model{
for(i in 1:length(rt)) {
rt[i] ~ dnorm(m, p) 
}
m ~ dnorm(0, 1/pow(500,2));
s ~ dunif(0,1000);
p <- 1/pow(s,2)
}"

###
### ...and corresponding "one-trick" normal model just for comparison.
###

trick_norm_model_str <- 
  "
model{
for (iter in 1:length(rt)){
p[iter] <- 1/(sqrt(2*3.141593) *s) * exp(-((rt[iter] - m)^2/(2 * s^2))) / 1000 ;
ones[iter] ~ dbern(p[iter]);
}
# Params to the distribution
m ~ dnorm(0, 1/pow(500,2));
s ~ dunif(0,1000);
}"

###
### A Three parameter Weibull distribution.
###

trick_weibull_model_str <- 
  "
model{
for (iter in 1:length(rt)){
p[iter] <- dweib(rt[iter] - shift ,shape, lambda) / 1000;

ones[iter] ~ dbern(p[iter]);
}

shape ~ dunif(0,20)
scale ~ dunif(0,500)
shift ~ dunif(0, 150)
lambda <- pow(1/(scale), shape)
}"


###
### Ex-gaussian distribution
###

trick_ex_gaussian_model_str <- 
  "
model{
for (iter in 1:length(rt)){
p[iter] <- l/2*exp(l/2*(2*m + l*s^2 - 2*rt[iter]))*2*phi(-(m+l*s^2-rt[iter])/(sqrt(2)*s)*sqrt(2)) / 1000 ;
ones[iter] ~ dbern(p[iter]);
}
# Params to the distribution
m ~ dunif(1, 1000);
s ~ dunif(1,max_s);
max_s <- 500
max_l <- 1/(max_s * 0.05)
l ~ dunif(0.000001, max_l)
}"

# Function generating samples from the posterior using JAGS.
run_rt_jags <- function(model_str, rt_data, variable_names) {
  model <- jags.model(textConnection(model_str), data = list(rt=rt_data, ones = rep(1, length(rt_data))), n.chains=3, n.adapt=1000)
  update(model, 1000) # Some burn in
  samples <- coda.samples(model, variable.names=variable_names, n.iter=5000)
  samples
}

# Running the normal model
norm_samples <- run_rt_jags(norm_model_str, rt_subject_48, c("m", "s"))

## Warning: Unused variable "ones" in data

## Compiling model graph
##    Resolving undeclared variables
##    Allocating nodes
##    Graph Size: 41
## 
## Initializing model

plot(norm_samples)

# Most likely point estimates of the parameters
norm_est <- colMeans(as.matrix(norm_samples))
round(norm_est, 4)

##      m      s 
## 217.21  53.92

# Plotting the best fitting normal distribution
hist(rt_subject_48, freq=F, xlim=c(0, 500), ylim=c(0, 0.012), breaks=6)
lines(0:500, dnorm(0:500, norm_est["m"], norm_est["s"]), col="red")

# Running the normal trick model
# Except for being slower it should give the same result as above
trick_norm_samples <- run_rt_jags(trick_norm_model_str, rt_subject_48, c("m", "s"))

## Compiling model graph
##    Resolving undeclared variables
##    Allocating nodes
##    Graph Size: 272
## 
## Initializing model

plot(trick_norm_samples)

# Most likely point estimates of the parameters
trick_norm_est <- colMeans(as.matrix(trick_norm_samples))
round(trick_norm_est, 4)

##      m      s 
## 217.11  53.72

# Plotting the best fitting normal distribution for the "trick" model
hist(rt_subject_48, freq=F, xlim=c(0, 500), ylim=c(0, 0.012), breaks=6)
lines(0:500, dnorm(0:500, trick_norm_est["m"], trick_norm_est["s"]), col="red")

# Running the three parameter Weibull "trick" model
weibull_samples <- run_rt_jags(trick_weibull_model_str, rt_subject_48, c("shape", "scale", "shift", "lambda"))

## Compiling model graph
##    Resolving undeclared variables
##    Allocating nodes
##    Graph Size: 155
## 
## Initializing model

plot(weibull_samples)

# Most likely point estimates of the parameters
weibull_est <- colMeans(as.matrix(weibull_samples))
round(weibull_est, 4)

##  lambda   scale   shape   shift 
##   0.001 114.125   2.058 117.577

# Plotting the best fitting Weibull distribution
hist(rt_subject_48, freq=F, xlim=c(0, 500), ylim=c(0, 0.012), breaks=6)
# Density function for the Three parameter Weibull
dweibull3 <- function(x, shape, scale, shift) {dweibull(x - shift, shape, scale)}
lines(0:500, dweibull3(0:500, weibull_est["shape"], weibull_est["scale"], weibull_est["shift"]), col="red")

# Running the Ex-Gaussian "trick" model
exgauss_samples <- run_rt_jags(trick_ex_gaussian_model_str, rt_subject_48, c("m", "s", "l"))

## Compiling model graph
##    Resolving undeclared variables
##    Allocating nodes
##    Graph Size: 387
## 
## Initializing model

plot(exgauss_samples)

# Most likely point estimates of the parameters
exgauss_est <- colMeans(as.matrix(exgauss_samples))
round(exgauss_est, 4)

##        l        m        s 
##   0.0225 169.0356  26.4306

# Plotting the best fitting Ex-Gaussian distribution
hist(rt_subject_48, freq=F, xlim=c(0, 500), ylim=c(0, 0.012), breaks=6)
# Density function for the Ex-Gaussian.
dexgaus <- function(x, m, s, l) {l/2*exp(l/2*(2*m + l*s^2 - 2*x))*2*pnorm(-(m+l*s^2-x)/(sqrt(2)*s)*sqrt(2))}
lines(0:500, dexgaus(0:500, exgauss_est["m"], exgauss_est["s"], exgauss_est["l"]), col="red")