#
# +---------------------------------------------------------+
# | Title: TransportFindings_StressValidation_Code.R
# +---------------------------------------------------------+

# Author: Dillon T. Fitch
# Last Updated: 7/28/2021

# Notes:
# (1) starts with summary data from a prior experiment (Fitch et al. (2020)) (data_for_modeling.RDS)
# (2) model survey responses as a multilevel ordered logit
# (3) summarize the relationships between HRV and survey responses (as predicted by the model) in 
#  a visual plot (figure 1)
# +---------------------------------------------------------+
# | LIBRARIES AND PACKAGES                                  |
# +---------------------------------------------------------+
library(brms)
library(ggplot2)
library(tidybayes)
library(stringr)
library(dplyr)
library(tidyr)

# -----------------------------------------------------------
# +---------------------------------------------------------+
# | Main Code                                               |
# +---------------------------------------------------------+
#
# Read data
d <- readRDS("data_for_modeling.RDS")

# data description--------------------------
# person_idx = index for unique person
# HF_HRV_sd = standard deviation of high frequency filtered (via MODWT transform) of RR intervals for
#             each road condition
# Condition = road condition (participants rode their bikes on5 roads with rests between each)
# speed_avg = average bicyclist speed (mps) for each road condition
# Vigilant_z = z-score of person-level vigilance
# Ability_z = z-score of person-level bicycling ability
# Desire_z = z-score of bicycling desire
# item = stress survey item (8 items (see Table 1 for item text))
# rating = item rating (5-point ordinal scale 1= Not at all, 2 = A little,
# 3 = Sometimes, 4 = Often, 5 = A lot)

# fit model (ensure 4 cores are available with excess RAM)-------------------
# estimated runtime: 16 minutes on AMD Ryzen 7 1800X eight-core processor
fit <- brm(rating ~ (Condition + HF_HRV_sd + HF_HRV_sd:speed_avg + HF_HRV_sd:Condition | person_idx) + 
                    (Condition + HF_HRV_sd + HF_HRV_sd:speed_avg + HF_HRV_sd:Condition | item ) +
                     Condition + HF_HRV_sd + HF_HRV_sd:speed_avg + HF_HRV_sd:Condition + 
                     Vigilant_z + Ability_z + Desire_z, 
           data=d, 
           family=cumulative("logit"), 
           iter = 2000,
           chains = 4,
           cores = 4,
           control = list(adapt_delta = 0.9, max_treedepth = 16),
           prior = c(set_prior("normal(0,2)", class = "b"), 
                     set_prior("student_t(6,0,1.5)", class = "Intercept"),
                     set_prior("student_t(6,0,1.5)", class = "sd"),
                     set_prior("lkj(2)", class = "cor"))
)  
saveRDS(fit,"fit.RDS")

# check that predictions are reasonable (not shown in paper)
pp_check(fit,type="bars",nsamples=100) # very good overall in-sample predictions
pp_check(fit,type="bars_grouped", group="Condition", nsamples=100) # moderate in-sample predictions by condition
pp_check(fit,type="bars_grouped", group="item", nsamples=100) # moderate in-sample predictions by item


# Evaluate results for main study objectives (HRV vs. survey responses) -------------------
# n= number of counterfactual conditions
n=length(unique(d$person_idx))*
  length(unique(d$item))*
  length(unique(d$Condition))*
  length(quantile(d$HF_HRV_sd,seq(0,1,by=.1)))

# counterfactual data
newdata <- data.frame(person_idx = rep(d$person_idx,each=11),
                      Condition = rep(d$Condition,each=11),
                      item = rep(d$item,each=11),
                      HF_HRV_sd = rep(quantile(d$HF_HRV_sd,seq(0,1,by=.1)),nrow(d)),
                      speed_avg = rep(mean(d$speed_avg),nrow(d)*11),
                      Vigilant_z = rep(0,nrow(d)*11),
                      Ability_z = rep(0,nrow(d)*11),
                      Desire_z = rep(0,nrow(d)*11))

# predict on counterfactual data (need lots of RAM, takes a few seconds)
pred <- predict(fit,newdata=newdata,summary=T)

# tidy data
d.plot <- cbind(newdata,pred)
d.plot <- d.plot %>% 
  pivot_longer(cols=starts_with("P("),
               names_to = "Class",
               values_to = "pk") %>%
  group_by(person_idx,item,Condition,HF_HRV_sd) %>%
  arrange(person_idx,item,Condition,HF_HRV_sd, Class) %>%
  mutate(`Cumulative Probability` = cumsum(pk))

d.plot$Class <- 
  d.plot$Class %>%
  recode_factor(`P(Y = 1)` = "Prob. = Not at all",
                `P(Y = 2)` = "Prob. \u2264 A little",
                `P(Y = 3)` = "Prob. \u2264 Sometimes",
                `P(Y = 4)` = "Prob. \u2264 Often",
                `P(Y = 5)` = "Prob. \u2264 A lot")

# summarize means for plotting
d.plot2 <- d.plot %>% ungroup() %>% 
  group_by(Condition,item,HF_HRV_sd,Class) %>% 
  summarize_if(is.numeric,mean)


item_names <- c(
  'Emotionally_Tired'="Emotionally\nTired",
  'Physically_Tired'="Physically\nTired",
  'Uncomfortable'="Uncomfortable",
  'Unprotect_Traffic'="Unprotected\nfrom traffic",
  'Unsafe'="Unsafe",
  'Worried_Turn_Cars'="Worried about\nturning cars",
  'Worry_Car_Doors'="Worried about\ncar doors",
  'Worry_Passing_Cars'="Worried about\npassing cars"
)

# generate and save Figure 1
jpeg(file="Figure1.jpg",width=6.5,height=7,units="in",res=900,pointsize = 8)
ggplot(d.plot,aes(y=`Cumulative Probability`,x=HF_HRV_sd, group=Class, color=Class))+
  geom_path(size=.1,alpha=.6)+
  geom_path(data=d.plot2,aes(y=`Cumulative Probability`,x=HF_HRV_sd, group=Class),
            size=1.1,alpha=1,color="black")+
  geom_path(data=d.plot2,aes(y=`Cumulative Probability`,x=HF_HRV_sd, group=Class, color=Class),
            size=.8,alpha=1)+
  facet_grid(item~Condition, labeller=labeller(item=item_names))+
  scale_colour_brewer(palette = "RdYlBu")+
  xlab("HF-HRV (std. dev. high-frequency RR intervals in milliseconds)")+
  ylab("Probability")+
  theme_dark()+
  theme(legend.position = "bottom",
        strip.text.y = element_text(angle=0),
        legend.title = element_blank())+
  guides(color = guide_legend(nrow = 2, byrow = TRUE))

dev.off()