###################################################
# R v4.0.5 code accompanying Rosenberger et al. (2021)
###################################################

library(bestNormalize) # for Yeo-Johnson transformation (v.1.7.0)
library(psych)         # for rotated principal component analysis (v.2.1.3)
library(lme4)          # for LMER fitting (v.1.1-26)
library(multcomp)      # for testing contrasts (v.1.4-16)

###################################################
# Part I: Read and prepare Data 
###################################################

# Path to the directory with the data files T1_animals.csv, T2_NA.csv, T3_NO.csv, T4_NH.csv, and T5_WH.csv
d_dir <- "path_to_dir" # Adjust !

# -------------------------------------------------

# Classes for columns to be read from T1_animals.csv
d_classes_animals <- c(Individual    = "character", 
                       DateOfBirth   = "NULL",
                       Taxonomy      = "NULL",
                       Breed         = "NULL",
                       SelectionLine = "factor", 
                       Pen           = "factor", 
                       Site          = "factor", 
                       Treatment     = "factor")

# Read T1_animals.csv to data.frame
animals_data <- read.table(file = file.path(d_dir, "T1_animals.csv"), 
                           sep = ",", header = TRUE, 
                           colClasses = d_classes_animals)

# Reorder levels in SelectionLine and Treatment columns
animals_data$SelectionLine <- factor(animals_data$SelectionLine, levels = c("Dwarf", "Dairy"))
animals_data$Treatment <- factor(animals_data$Treatment, levels = c("COG", "POS", "ISO"))

# -------------------------------------------------

# Read T2_NA.csv to data.frame
NA_data <- read.table(file = file.path(d_dir, "T2_NA.csv"), 
                       sep = ",", header = TRUE)

# Add animal information to NA_data
NA_data <- merge(NA_data, animals_data, by = "Individual", all.x = TRUE)

# Change data type of Individual to factor
NA_data$Individual <- as.factor(NA_data$Individual)

# -------------------------------------------------

# Read T3_NO.csv to data.frame
NO_data <- read.table(file = file.path(d_dir, "T3_NO.csv"), 
                      sep = ",", header = TRUE)

# Add animal information to NO_data
NO_data <- merge(NO_data, animals_data, by = "Individual", all.x = TRUE)

# Change data type of Individual to factor
NO_data$Individual <- as.factor(NO_data$Individual)

# -------------------------------------------------

# Read T4_NH.csv to data.frame
NH_data <- read.table(file = file.path(d_dir, "T4_NH.csv"), 
                      sep = ",", header = TRUE)

# Add animal information to NH_data
NH_data <- merge(NH_data, animals_data, by = "Individual", all.x = TRUE)

# Change data type of Individual to factor
NH_data$Individual <- as.factor(NH_data$Individual)

# -------------------------------------------------

# Read T5_WH.csv to data.frame
WH_data <- read.table(file = file.path(d_dir, "T5_WH.csv"), 
                      sep = ",", header = TRUE)

# Add animal information to WH_data
WH_data <- merge(WH_data, animals_data, by = "Individual", all.x = TRUE)

# Change data type of Individual to factor
WH_data$Individual <- as.factor(WH_data$Individual)

###################################################
# Part II: Principal component analysis (PCA)
###################################################

# Define function for Yeo-Johnson transformation and PCA
make_pca <- function(pca_data, vars, nfac = 2) {
  pca_data <- pca_data[, colnames(pca_data) %in% vars]
  pca_data <- as.data.frame(lapply(pca_data, function(i) yeojohnson(i)$x.t))
  pca <- principal(pca_data, rotate = "varimax", nfactors = nfac, covar = FALSE)
  return(pca)
}

# Selected variables (selection as explained in the paper)
NA_vars <- c("bs_HeartRate", "Inactive_Duration", "Vocalising_Frequency", "SegmentChange_Frequency", "InnerSegments_Duration")

NO_vars <- c("bs_HeartRate", "Inactive_Duration", "ObjectContact_Frequency", "SegmentChanges_Frequency")

NH_vars <- c("bs_HeartRate", "Inactive_Duration", "Vocalising_Frequency", "HumanContact_Frequency", "SegmentChanges_Frequency") 

WH_vars <- c("bs_HeartRate", "WeighingScore", "ExitingScore")

# Fit PCAs
NA_pca <- make_pca(pca_dat = NA_data, vars = NA_vars, nfac = 2) 

NO_pca <- make_pca(pca_dat = NO_data, vars = NO_vars, nfac = 2) 

NH_pca <- make_pca(pca_dat = NH_data, vars = NH_vars, nfac = 2) 

WH_pca <- make_pca(pca_dat = WH_data, vars = WH_vars, nfac = 1) 

# Inspect PCAs 

print(NA_pca)

print(NO_pca)

print(NH_pca) 

print(WH_pca)

###################################################
# Part III: LMER with rotated components (and Entering score in WH)
###################################################
  
# Fit LMER models for rotated components 
NA_RC1_model <- lmer(RC1 ~ 0 + Treatment:SelectionLine + (1 |Site/Pen), data = cbind(NA_data, NA_pca$scores))
NA_RC2_model <- lmer(RC2 ~ 0 + Treatment:SelectionLine + (1 |Site/Pen), data = cbind(NA_data, NA_pca$scores))

NO_RC1_model <- lmer(RC1 ~ 0 + Treatment:SelectionLine + (1 |Site/Pen), data = cbind(NO_data, NO_pca$scores))
NO_RC2_model <- lmer(RC2 ~ 0 + Treatment:SelectionLine + (1 |Site/Pen), data = cbind(NO_data, NO_pca$scores))

NH_RC1_model <- lmer(RC1 ~ 0 + Treatment:SelectionLine + (1 |Site/Pen), data = cbind(NH_data, NH_pca$scores))
NH_negRC2_model <- lmer(-RC2 ~ 0 + Treatment:SelectionLine + (1 |Site/Pen), data = cbind(NH_data, NH_pca$scores)) 

WH_PC1_model <- lmer(PC1 ~ 0 + Treatment:SelectionLine + (1 |Site/Pen), data = cbind(WH_data, WH_pca$scores))

# Fit LMER models for entering score 

WH_data$EnteringScore_transformed <- yeojohnson(WH_data$EnteringScore)$x.t

WH_Entering_model <- lmer(EnteringScore_transformed ~ 0 + Treatment:SelectionLine + (1 |Site/Pen), data = WH_data)

# Inspect models and calculate p-values for fixed effect
summary(NA_RC1_model)
summary(glht(NA_RC1_model), test = adjusted('none'))

summary(NA_RC2_model) 
summary(glht(NA_RC2_model), test = adjusted('none'))

summary(NO_RC1_model) 
summary(glht(NO_RC1_model), test = adjusted('none'))

summary(NO_RC2_model) 
summary(glht(NO_RC2_model), test = adjusted('none'))

summary(NH_RC1_model) 
summary(glht(NH_RC1_model), test = adjusted('none'))

summary(NH_negRC2_model) 
summary(glht(NH_negRC2_model), test = adjusted('none'))

summary(WH_PC1_model) 
summary(glht(WH_PC1_model), test = adjusted('none'))

summary(WH_Entering_model) 
summary(glht(WH_Entering_model), test = adjusted('none'))
# -------------------------------------------------
  
# Define contrast matrix for Treatment and Selection line contrasts
C <- rbind(
  # Treatment contrasts
  "TreatmentPOS:SelectionLineDwarf - TreatmentCOG:SelectionLineDwarf" = c(-1, 1, 0, 0, 0, 0),
  "TreatmentISO:SelectionLineDwarf - TreatmentPOS:SelectionLineDwarf" = c(0, -1, 1, 0, 0, 0),
  "TreatmentISO:SelectionLineDwarf - TreatmentCOG:SelectionLineDwarf" = c(-1, 0, 1, 0, 0, 0),
  "TreatmentPOS:SelectionLineDairy - TreatmentCOG:SelectionLineDairy" = c(0, 0, 0, -1, 1, 0),
  "TreatmentISO:SelectionLineDairy - TreatmentPOS:SelectionLineDairy" = c(0, 0, 0, 0, -1, 1),
  "TreatmentISO:SelectionLineDairy - TreatmentCOG:SelectionLineDairy" = c(0, 0, 0, -1, 0, 1),
  # Selection line contrasts
  "TreatmentCOG:SelectionLineDairy - TreatmentCOG:SelectionLineDwarf" = c(-1, 0, 0, 1, 0, 0),
  "TreatmentPOS:SelectionLineDairy - TreatmentPOS:SelectionLineDwarf" = c(0, -1, 0, 0, 1, 0),
  "TreatmentISO:SelectionLineDairy - TreatmentISO:SelectionLineDwarf" = c(0, 0, -1, 0, 0, 1)
)
  
# test contrasts
summary(glht(NA_RC1_model, linfct = C), test = adjusted('none'))
summary(glht(NA_RC2_model, linfct = C), test = adjusted('none'))

summary(glht(NO_RC1_model, linfct = C), test = adjusted('none'))
summary(glht(NO_RC2_model, linfct = C), test = adjusted('none'))

summary(glht(NH_RC1_model, linfct = C), test = adjusted('none'))
summary(glht(NH_negRC2_model, linfct = C), test = adjusted('none'))
                
summary(glht(WH_PC1_model, linfct = C), test = adjusted('none'))
summary(glht(WH_Entering_model, linfct = C), test = adjusted('none'))
# -------------------------------------------------

# define function to calculate 95% CI for fixed effect estimates
calc_estimate_CI <- function(model) { 
  
  # Define grid to calculate fitted values plus confidence band
  prediction_grid <- expand.grid(Treatment = factor(c("COG", "POS", "ISO"), levels = c ("COG", "POS", "ISO")), 
                                 SelectionLine = factor(c("Dwarf", "Dairy"), levels = c("Dwarf", "Dairy")))
  
  # Calculate fitted values (considering fixed effects only)
  prediction_grid$fitted <- predict(model, newdata = prediction_grid, re.form = ~ 0)  
  
  # Parametric bootstrapping (considering fixed effects only)
  boot_predict <- bootMer(model, nsim = 10000, seed = 1,
                          FUN = function(x) predict(x, newdata = prediction_grid, re.form = ~ 0),  
                          parallel = "multicore", 
                          ncpus = 4) 
  
  # Determine confidence band from bootstrap results
  CI <- apply(boot_predict$t, 2, quantile, probs = c(0.025, 0.975))
  
  prediction_grid <- cbind(prediction_grid, t(CI))
  
  return(prediction_grid)
  
}

# Calculate CIs 
CI_NA_RC1  <- calc_estimate_CI(model = NA_RC1_model)
CI_NA_RC2  <- calc_estimate_CI(model = NA_RC2_model)

CI_NO_RC1  <- calc_estimate_CI(model = NO_RC1_model)
CI_NO_RC2  <- calc_estimate_CI(model = NO_RC2_model)

CI_NH_RC1  <- calc_estimate_CI(model = NH_RC1_model)
CI_NH_negRC2  <- calc_estimate_CI(model = NH_negRC2_model)

CI_WH_PC1   <- calc_estimate_CI(model = WH_PC1_model)
CI_WH_Entering <- calc_estimate_CI(model = WH_Entering_model)