Estimating the mean phenotype for each DGRP line

load_and_tidy_fitness_data <- function(){
  female_fitness <- read_csv("data/input/female_fitness.csv") %>% 
    arrange(line) %>%
    split(paste(.$block, .$line)) %>%
    map(~ .x %>% mutate(replicate = 1:n())) %>%
    bind_rows() %>%
    mutate(vial = 1:n(),
           vial = paste("F", line, block, replicate, sep = "_"),
           line = paste("line_", line, sep = "")) %>% 
    dplyr::select(-num.F.late1, -num.F.late2, -num.laying.F.early, -replicate) %>%
    mutate(female.fitness.early = replace(female.fitness.early, is.na(female.fitness.early), 0),
           female.fitness.late = replace(female.fitness.late, is.na(female.fitness.late), 0)) 
  
  male_fitness <- read_csv("data/input/male_fitness.csv") %>% 
    arrange(line) %>%
    split(paste(.$block, .$line)) %>%
    map(~ .x %>% mutate(replicate = 1:n())) %>%
    bind_rows() %>%
    mutate(vial = paste("M", line, block, replicate, sep = "_"),
           line = paste("line_", line, sep = "")) %>%
    dplyr::select(-num.DGRP.males, -replicate)
  
  # Remove the small number of observations where the females had no offspring, so we cannot measure the
  # rival vs focal male competitive abilities
  # NOTE: Not needed any more, because now I impute this as missing data
  # exclude_early <- male_fitness$early.male.rival + male_fitness$early.male.focal == 0
  # exclude_late <- male_fitness$late.male.rival + male_fitness$late.male.focal == 0
  # male_fitness$early.male.rival[exclude_early] <- NA
  # male_fitness$early.male.focal[exclude_early] <- NA
  # male_fitness$late.male.rival[exclude_late] <- NA
  # male_fitness$late.male.focal[exclude_late] <- NA
  
  # Save these tidy files, for data archiving
  write.csv(female_fitness, file = "data/input/female_fitness_CLEANED.csv", row.names = FALSE)
  write.csv(male_fitness, file = "data/input/male_fitness_CLEANED.csv", row.names = FALSE)
  
  list(female_fitness = female_fitness, male_fitness = male_fitness)
}


# This function calculates the traditional line means (i.e. simple averages) for each line
get_line_means <- function(){
  left_join(
    load_and_tidy_fitness_data()[[1]] %>% 
      select(line, female.fitness.early, female.fitness.late) %>% 
      group_by(line) %>%
      summarise(female.fitness.early = mean(female.fitness.early), 
                female.fitness.late = mean(female.fitness.late)),
    load_and_tidy_fitness_data()[[2]] %>% 
      group_by(line) %>%
      summarise(early.male.rival = sum(early.male.rival),
                early.male.focal = sum(early.male.focal),
                late.male.rival = sum(late.male.rival),
                late.male.focal = sum(late.male.focal),
                male.fitness.early = early.male.focal / (early.male.focal + early.male.rival),
                male.fitness.late = late.male.focal / (late.male.focal + late.male.rival)) %>%
      select(line, male.fitness.early, male.fitness.late),
    by = "line") %>% as.data.frame()
}

# A function to calculate each line's expected fitness, using brms models to adjust for block + vial random effects
# For the male late life assay, we also adjust for the number of competitor males that died 
# (this assumes that their deaths happened randomly with respect to which line we were measuring)
get_predicted_line_means <- function(overwrite = FALSE){
  
  # If it's already calculated, just load up the file. Otherwise, make the file
  out.file <- "data/derived/predicted_line_means.csv"
  if(file.exists(out.file) & !overwrite) return(read_csv(out.file))
  
  ####### Load the individual-level data for females
  female_fitness <- load_and_tidy_fitness_data()[[1]] 
  
  ####### Load the individual-level data for males, and impute missing values
  
  # For one vial, there is no early-life measurement because the females did not produce offspring. 
  # Therefore we impute that datapoint and make 5 imputed datasets using the package 'mice'
  # See this link for explanation: https://cran.r-project.org/web/packages/brms/vignettes/brms_missings.html
  male_fitness <- load_and_tidy_fitness_data()[[2]]
  male_fitness <- male_fitness %>% 
    mutate(early.male.rival = replace(early.male.rival, early.male.rival + early.male.focal == 0, NA),
           early.male.focal = replace(early.male.focal, is.na(early.male.rival), NA)) 
  male_fitness_imputed <- mice(male_fitness, m = 5, print = FALSE, seed = 1)
  male_fitness_imputed <- lapply(1:5, function(i) { # for each imputed dataset, 
    complete(male_fitness_imputed, i) %>% 
      mutate(total1 = early.male.rival + early.male.focal, # tally up the focal+rival offspring for binomial model
             total2 = late.male.rival + late.male.focal,
             num.GFP.males = as.numeric(scale(num.GFP.males))) %>% # scale this covariate
      # For late-life assays where the 5 focal males all died before reaching age 14, 
      # we need to assign them zero fitness. We did this by assuming that the rivals sired 100/100 offspring.
      # This scenario happened in 47 vials (out of 810)
      mutate(late.male.focal = replace(late.male.focal, total2 == 0, 100),
             total2 = replace(total2, total2 == 0, 100))
  })

  ####### Run model on the female data 
  
  prior_females <- 
    c(set_prior("cauchy(0, 0.1)", class = "sd", resp = 'femalefitnessearly', group = "line"),
      set_prior("cauchy(0, 0.1)", class = "sd", resp = 'femalefitnessearly', group = "block"),
      set_prior("cauchy(0, 0.1)", class = "sd", resp = 'femalefitnessearly', group = "vial"),
      set_prior("cauchy(0, 0.1)", class = "sd", resp = 'femalefitnesslate', group = "line"),
      set_prior("cauchy(0, 0.1)", class = "sd", resp = 'femalefitnesslate', group = "block"),
      set_prior("cauchy(0, 0.1)", class = "sd", resp = 'femalefitnesslate', group = "vial"))
  
  prior_males <- 
    c(set_prior("cauchy(0, 0.1)", class = "sd", resp = 'earlymalefocal', group = "line"),
      set_prior("cauchy(0, 0.1)", class = "sd", resp = 'earlymalefocal', group = "block"),
      set_prior("cauchy(0, 0.1)", class = "sd", resp = 'earlymalefocal', group = "vial"),
      set_prior("cauchy(0, 0.1)", class = "sd", resp = 'latemalefocal', group = "line"),
      set_prior("cauchy(0, 0.1)", class = "sd", resp = 'latemalefocal', group = "block"),
      set_prior("cauchy(0, 0.1)", class = "sd", resp = 'latemalefocal', group = "vial"))
      
   
  
  
  female1 <- bf(female.fitness.early ~ (1 | p | line) + (1 | q | block) + (1 | r | vial))
  female2 <- bf(female.fitness.late ~  (1 | p | line) + (1 | q | block) + (1 | r | vial))
  female_model <- brm(female1 + female2, 
                      family = "poisson", data = female_fitness,
                      prior = prior_females,
                      iter = 5000, chains = 4, cores = 4,
                      seed = 1, control = list(adapt_delta = 0.999, max_treedepth = 15)) 
  writeLines(capture.output(summary(female_model)), con = "data/derived/model_summary_females.txt")
  
  saveRDS(female_model, "data/derived/female_fitness_prediction_model.rds") # maybe delete later - probably not needed?
  
  
  ### Run model on male data 
  # Fit num.GFP.males, but only for the late life data
  male1 <- bf(early.male.focal | trials(total1) ~ (1 | p | line) + (1 | q | block) + (1 | r | vial))
  male2 <- bf(late.male.focal  | trials(total2) ~ num.GFP.males + (1 | p | line) + (1 | q | block) + (1 | r | vial))
  
  # Run 5 models, on each of the 5 imputed datasets, and combine the results
  male_model <- brm_multiple(male1 + male2,
                             family = "binomial", 
                             prior = prior_males,
                             data = male_fitness_imputed,
                             iter = 5000, chains = 4, cores = 4,
                             seed = 1, control = list(adapt_delta = 0.999, max_treedepth = 15)) 
  writeLines(capture.output(summary(male_model)), con = "data/derived/model_summary_males.txt")
  saveRDS(male_model, "data/derived/male_fitness_prediction_model.rds") # maybe delete later - probably not needed?
  male_model <- readRDS("data/derived/male_fitness_prediction_model.rds") 
  
  ####### Define new data for prediction
  new.data.female <- female_fitness %>% 
    select(line) %>% distinct() %>% arrange(line)
  
  new.data.male <- male_fitness %>% 
    select(line) %>% distinct() %>% arrange(line) %>% 
    mutate(total1 = 100, # I checked, and changing this number does not affect the line means, as expected
           total2 = 100, 
           num.GFP.males = male_fitness %>% .$num.GFP.males %>% mean(na.rm=TRUE))
  
  # Get predicted lines means for early and late life *female* fitness on linear predictor scale (i.e. fairly normal)
  female_line_preds <- data.frame(new.data.female, 
                                  fitted(female_model, 
                                         newdata = new.data.female, 
                                         re_formula = "~ ( 1 | line)", 
                                         scale = "linear")) %>%
    select(line, Estimate.femalefitnessearly, Estimate.femalefitnesslate) %>% 
    rename(female.fitness.early = Estimate.femalefitnessearly,
           female.fitness.late = Estimate.femalefitnesslate) %>%
    mutate(female.fitness.early = as.numeric(scale(female.fitness.early)), # scale the line means
           female.fitness.late = as.numeric(scale(female.fitness.late)))
  
  # Get predicted lines means for early and late life *male* fitness on linear predictor scale (i.e. fairly normal)
  male_line_preds <- data.frame(new.data.male, 
                                fitted(male_model, 
                                       newdata = new.data.male, 
                                       re_formula = "~ (1 | line)", 
                                       scale = "linear")) %>%
    select(line, Estimate.earlymalefocal, Estimate.latemalefocal) %>% 
    rename(male.fitness.early = Estimate.earlymalefocal,
           male.fitness.late = Estimate.latemalefocal) %>%
    mutate(male.fitness.early = as.numeric(scale(male.fitness.early)), # scale the line means
           male.fitness.late = as.numeric(scale(male.fitness.late))) 

  ## Define a function to get the ICC for line, for each posterior sample of predicted line means
  get_ICC_line <- function(predictions){
    cols <- str_detect(names(predictions), "X")
    first_posterior_column <- which(cols)[1] 
    n_draws <- sum(cols)
    plan("multiprocess")
    future_lapply(0:(n_draws - 1), function(i){
       ICCbare(predictions$line, predictions[, i + first_posterior_column])
    }) %>% unlist()
  }
  
  # Calculate ICC for each sample, separately for old and young females
  # For females, this is done on the estimated line mean progeny produced (i.e. response scale)
  predicted <- fitted(female_model, summary = FALSE)
  Line_ICC_female_early <- get_ICC_line(data.frame(female_fitness, t(predicted[, , 1])))
  Line_ICC_female_late <- get_ICC_line(data.frame(female_fitness, t(predicted[, , 2])))
  
  # Do the same for males
  # For males, this is done on the estimated proportion progeny sired (i.e. response scale)
  new_data <- male_fitness %>% 
    select(block, line, vial.ID) %>% 
    mutate(total1 = 100, 
           total2 = 100, 
           num.GFP.males = male_fitness %>% .$num.GFP.males %>% mean(na.rm=TRUE))
  predicted <- fitted(male_model, summary = FALSE, newdata = new_data)
  Line_ICC_male_early <- get_ICC_line(data.frame(new_data, t(predicted[, , 1])))
  Line_ICC_male_late <- get_ICC_line(data.frame(new_data, t(predicted[, , 2])))
  
  ICC <- list(Line_ICC_female_early, 
              Line_ICC_female_late,
              Line_ICC_male_early,
              Line_ICC_male_late) %>% lapply(posterior_summary) %>% 
    do.call("rbind", .) %>% as_tibble() %>%
    mutate(Fitness_trait = c("Female early-life fitness", "Female late-life fitness",
                             "Male early-life fitness", "Male late-life fitness"))
  write.csv(ICC, file = "data/derived/ICC.csv", row.names = TRUE)
  

  # Record which block(s) each line was measured in. Generally one block, except for ref line (all blocks)
  line_blocks <- rbind(female_fitness %>% select(block, line), 
                       male_fitness   %>% select(block, line)) %>% 
    mutate(block = replace(block, line == "352", "reference_line")) %>% 
    distinct() %>% 
    group_by(line) %>% summarise(block = paste0(block, collapse = ", ")) 
  
  predicted_line_means <- female_line_preds %>%
    left_join(male_line_preds, by = "line") %>% 
    left_join(line_blocks, by = "line") %>%
    arrange(line) 
  
  write.csv(predicted_line_means, out.file, row.names = FALSE)
}

line_means <- get_line_means()
predicted_line_means <- get_predicted_line_means(overwrite = FALSE) 

cat(readLines('data/derived/model_summary_females.txt'), sep = '\n')

cat(readLines('data/derived/model_summary_males.txt'), sep = '\n')

read.csv('data/derived/ICC.csv') %>% 
  select(Fitness_trait, Estimate, Est.Error, Q2.5, Q97.5) %>%
  rename(`Fitness trait` = Fitness_trait, `Est. Error` = Est.Error, `Lower 95% CI` = Q2.5, `Upper 95% CI` = Q97.5) %>% 
  mutate(across(where(is.numeric), ~ .x * 100)) %>%
  kable(digits = 2) %>% kable_styling(full_width = FALSE)