##############Joint work with Juxin Liu, Yang Liu and Peng Zhang#############
library(rstan) 
library(bayesplot)
library(xtable)
library(parallel)
library(dplyr)
#####################################Data####################################
source("data.R") # set the working directory to where the file is stored
load("data/collec.RData")
Time <- c(rep(0, 83*7), rep(1, 83*7))
x_all <- as.matrix(cbind(rep(1,83*14), scale(xx1), scale(xx2), Time))
x_R <- rep(rate, each=83)
#############################################################################


##########################Parameters for Beta Prior##########################
## parameters for beta prior
beta2 <- function(para) {
  para[3] <- 0.025
  para[4] <- 0.975
  sum.sqr <- function(x) {
    sum1 <- (pbeta(para[1], shape1 = x[1], shape2 = x[2]) - para[3])^2
    sum2 <- (pbeta(para[2], shape1 = x[1], shape2 = x[2]) - para[4])^2
    ans <- sum1 + sum2
  }
  opt <- optim(c(1, 1), sum.sqr, lower = c(0.01, 0.01), method = "L-BFGS-B")
  #opt<-optim(c(1,1),sum.sqr)
  list(a = opt$par[1], b = opt$par[2])
}
## hyper prior parameters for prior 1
sbeta <- beta2(c(0.12,0.28)) 
## hyper prior parameters for prior 2
#sbeta <- beta2(c(0.4,0.6))
## hyper prior parameters for prior 3
#sbeta <- beta2(c(0.7,0.89))
gamma1 <- sbeta$a; gamma2 <- sbeta$b
#############################################################################


###############################Naive Model###################################
inits1 <- list(beta_m = rep(1, 4), phi = rep(0.1, 83), 
               sd_tau = 0.5, alpha = 0.2) 
inits2 <-list( beta_m = rep(-0.1, 4), phi = rep(1, 83), 
               sd_tau = 1, alpha = 0.6)
inits <- list(inits1,inits2)
stanfit_naive <-stan("naive.stan", data = list(n = 83, N = 1162, y = yy, 
                                                    pop_tn = pop_tn, K = 4, 
                                                    x_all = x_all, x_R = x_R, 
                                                    W = W, W_n = W_n, 
                                                    gamma1 = gamma1, 
                                                    gamma2 = gamma2), 
                     chains = 2, warmup = 3000, iter = 6000, 
                     save_warmup = FALSE, init = inits, seed=1234,
                     control = list(adapt_delta = 0.95, max_treedepth = 15))
## output
print(stanfit_naive, digits=2,pars=c( "beta_m", "tau","alpha"), 
      probs=c(0.025, 0.975)) 
traceplot(stanfit_naive, pars=c("beta_m", "tau","alpha"))
pairs(stanfit_naive, pars=c("beta_m","tau","alpha"))
#############################################################################


##################################ME Model###################################
inits1 <- list(beta_m = rep(1, 4), beta_R = rep(1, 1),
               phi = rep(0.1, 83), tau = 0.5, alpha = 0.2) 
inits2 <-list( beta_m = rep(-0.1, 4), beta_R = rep(1, 1),
               phi = rep(1, 83), tau = 1, alpha = 0.6)
inits <- list(inits1,inits2)
stanfit_ME <-stan("ME.stan", data = list(n = 83, N = 1162, y = yy, 
                                                  pop_tn = pop_tn, K = 4, 
                                                  x_all = x_all, x_R = x_R,
                                                  W = W, W_n = W_n, 
                                                  gamma1 = gamma1, 
                                                  gamma2 = gamma2), 
                    chains = 2, warmup = 3000, iter = 6000,
                    save_warmup = FALSE, init = inits, seed = 1234,
                    control = list(adapt_delta = 0.95, max_treedepth = 15))
## output
print(stanfit_ME, digits=2,pars=c( "beta_m", "tau","alpha"), 
      probs=c(0.025, 0.975)) 
traceplot(stanfit_ME,pars=c("beta_m", "tau","alpha"))
pairs(stanfit_ME, pars=c("beta_m","tau","alpha"))
#############################################################################


#################################Trace plots-NAIVE phi#################################
parse_exp <- function(s) {
  parse(text = as.expression(s))
}

var_names <- list(
  "beta_m[1]" = parse_exp("beta[0]"), 
  "beta_m[2]" = parse_exp("beta[1]"), 
  "beta_m[3]"= parse_exp("beta[2]"), 
  "beta_m[4]" = parse_exp("beta[3]"), 
  "tau" = parse_exp("tau[phi]"), 
  "alpha" = parse_exp("alpha[phi]"))

var_labeller <- function(variable,value){
  return(var_names[value])
}
color_scheme_set("brightblue")
mcmc_trace(as.array(stanfit_naive), 
           pars = c("beta_m[1]", "beta_m[2]", "beta_m[3]", 
                    "beta_m[4]", "tau", "alpha"), 
           facet_args = list(labeller = var_labeller)) + theme_bw()
#############################################################################


#################################Trace plots-ME eta#################################
parse_exp <- function(s) {
  parse(text = as.expression(s))
}

var_names <- list(
  "beta_m[1]" = parse_exp("beta[0]"), 
  "beta_m[2]" = parse_exp("beta[1]"), 
  "beta_m[3]"= parse_exp("beta[2]"), 
  "beta_m[4]" = parse_exp("beta[3]"), 
  "tau" = parse_exp("tau[eta]"), 
  "alpha" = parse_exp("alpha[eta]"))

var_labeller <- function(variable,value){
  return(var_names[value])
}
color_scheme_set("brightblue")
mcmc_trace(as.array(stanfit_ME), 
           pars = c("beta_m[1]", "beta_m[2]", "beta_m[3]", 
                    "beta_m[4]", "tau", "alpha"), 
           facet_args = list(labeller = var_labeller)) + theme_bw()
#############################################################################




###################################DHARMa####################################
library(DHARMa)
y_rep <- extract(stanfit_naive)$y_rep
y_rep_me <- extract(stanfit_ME)$y_rep
res_me <- createDHARMa(simulatedResponse = t(y_rep_me), observedResponse = yy, 
                   fittedPredictedResponse = apply(t(y_rep_me), 1, median), 
                   integerResponse = T)
res_naive <- createDHARMa(simulatedResponse = t(y_rep), observedResponse = yy, 
                      fittedPredictedResponse	= apply(t(y_rep), 1, median), 
                      integerResponse = T)
testUniformity(res_naive); testUniformity(res_me)
#############################################################################


#############################################################################
#########Plots for yy, state level teen driver & teen population#############
## plot teenage car crash counts
qplot(yy, bins=17, geom = "histogram", fill = I("darkblue"), colour = I("grey")) + 
  ylab("Frequency") + xlab("Teen-driver fatal car crash counts") + theme_gray()

## Plot teen driver (state level)
before <- ggplot(teen_d_plot, aes(year, under19)) + geom_line() +
  xlab("Year") + ylab("Number of licensed drivers")+
  geom_segment(aes(x = 1997, y = 432000, xend = 1997, yend = 370000),
               arrow = arrow(length = unit(0.5, "cm")), colour = "red")+
  annotate("text", x = 1999.3, y = 400000, label = "Graduated Drivers Licence", 
           colour = "red") + #ylim(370000, 431000) +
  scale_x_continuous(labels = teen_d_plot$year, breaks = teen_d_plot$year) + 
  theme_gray()  

## plot teenage population (state level)
pop_state_plot <- data.frame(year=1990:2004, popT=pop_tn_state_plot)
after <- ggplot(pop_state_plot, aes(year, popT)) + geom_line() +
  xlab("Year") + ylab("Number of teenager population")+
  geom_segment(aes(x = 1997, y = 730000, xend = 1997, yend = 660000),
               arrow = arrow(length = unit(0.5, "cm")), colour = "red")+
  annotate("text", x = 1999.3, y = 690000, label = "Graduated Drivers Licence", 
           colour = "red")+
  scale_x_continuous(labels = pop_state_plot$year, 
                     breaks = pop_state_plot$year) + theme_gray()

##############################Arrange Plots##################################   
library(cowplot)
plot_grid( after, before,ncol = 1)
#############################################################################






