# The ventral hippocampus-medial prefrontal pathway is swayed by trait anxiety for slow, but not fast, escape decisions
# Final manuscript code
# Bowen J Fung, 2018
# bjfung@caltech.edu

# Load required packages
require(data.table)
require(lme4)
require(stargazer)
require(coxme)
require(ez)

# Read data
dt <- fread("originalData.csv", drop = 1)

# Reclass variables
dt <- within(dt, {
  subject <- as.factor(subject)
  flight <- as.numeric(flight)
  reward <- as.numeric(reward)
  color <- as.factor(color)
  reward_level <- as.factor(reward_level)
  shock_level <- as.factor(shock_level)
  shockornot <- as.factor(shockornot)
  trial <- as.integer(trial)
  escape <- as.numeric(escape) # needs to be numeric for survival analysis
  anxiety <- as.numeric(anxiety)
  attack <- as.numeric(attack)
  age <- as.numeric(age)
  gender <- as.factor(gender)
  BAS_drive <- as.numeric(BAS_drive)
  BAS_fun <- as.numeric(BAS_fun)
  BAS_reward <- as.numeric(BAS_reward)
  BIS <- as.numeric(BIS)
})

# Remove 17 possibly

# Remove control condition
dt <- droplevels(dt[color != 4])

# Rename predator conditions
setattr(dt$color,"levels",c("Fast", "Medium", "Slow"))

# Trait anxiety and flight initiation distance mixed model (use lmerTest for t-vals and dfs)
lFit <- lmer(flight ~ color*anxiety + (1|subject), dt[escape == 1])
stargazer(lFit, header = FALSE, ci = TRUE)

# BIS and flight initiation distance
lFit <- lmer(flight ~ color*anxiety*BIS + (1|subject), dt[escape == 1])
stargazer(lFit, header = FALSE, ci = TRUE)

## Test for equality of variances in empirical attack times (i.e. attack times in experimental design)
temp <- dt[subject == 1, .(attack = attack), by = color] # take first subject as example. note that dfs reflect imbalanced conditions
var.test(attack ~ color, data = temp[color %in% c("Fast","Medium")])
var.test(attack ~ color, data = temp[color %in% c("Medium","Slow")])
var.test(attack ~ color, data = temp[color %in% c("Fast","Slow")])
rm(temp)

## Test for equality of variances in median FID
temp <- dt[escape == 1, .(fid = median(flight)), by = .(subject, color)]
var.test(fid ~ color, data = temp[color %in% c("Fast","Medium")])
var.test(fid ~ color, data = temp[color %in% c("Medium","Slow")])
var.test(fid ~ color, data = temp[color %in% c("Fast","Slow")])
rm(temp)

## Test for differences in variances of individual subjects, across subjects
temp <- dt[escape == 1, .(fidVar = var(flight)), by = .(subject, color)]
t.test(fidVar ~ color, data = temp[color %in% c("Fast","Medium")])
t.test(fidVar ~ color, data = temp[color %in% c("Medium","Slow")])
t.test(fidVar ~ color, data = temp[color %in% c("Fast","Slow")])
temp[,mean(fidVar), by = color] # mean variances
rm(temp)

## Test for equality of variances in pooled data
temp <- dt[escape == 1, .(fid = median(flight), traitAnx = mean(anxiety)), by = .(subject, color)]
temp <- temp[, fast := ifelse(color %in% c("Fast","Medium"), "Fast", "Slow")]
var.test(fid ~ fast, data = temp)
rm(temp)

## Mixed model on pooled data
temp <- dt[, fast := ifelse(color %in% c("Fast","Medium"), "Fast", "Slow")]
lFit <- lmer(flight ~ fast*anxiety + (1|subject), temp[escape == 1])
stargazer(lFit, header = FALSE, ci = TRUE)
rm(temp)

# Survival analysis
dt[, proximity := (max(flight) - flight)/max(flight)*100] # add proximity variable
cFit <- coxme(Surv(proximity,escape) ~ color*anxiety + (1 | subject), data = dt)
summary(cFit)

# Trait anxiety and performance
temp <- dt[, .(traitAnx = mean(anxiety), performance = sum(reward),
               esc = sum(as.numeric(as.character(escape))) / length(escape)), by = .(subject, color)]
temp[, normPerformance := (performance - mean(performance)) / sd(performance), by = color] # z-score for performance

## Test trait anxiety and economic performance
ezANOVA(temp, dv = normPerformance, wid = subject, within = .(color), between = .(traitAnx), type = 3)

## Test correlation between economic performance and escape success
temp2 <- temp[, .(normPerformance = mean(normPerformance), esc = mean(esc)), by = subject]
cor.test(temp2$normPerformance, temp2$esc, method = "pearson")

## Test trait anxiety and escape success
ezANOVA(temp, dv = esc, wid = subject, within = .(color), between = .(traitAnx), type = 3)
ezANOVA(subset(temp, color == "Slow"), dv = esc, wid = subject, between = .(traitAnx), type = 3) # simple effects within slow predator condition