hell yes

library(dplyr)
library(doBy)
library(lubridate)
library(ggplot2)
library(ggforce)
library(tidyverse)
library(lme4)
library(nlme)
library(equatiomatic)
library(sjPlot)
library(ggeffects)
library(LMERConvenienceFunctions)

age <- function(birth, base = Sys.Date()){
  i <- interval(birth, base)
  p <- as.period(i)
  year(p)
}

setupDataframes <- function(df, geneStat) {
  df <- order_by(~usubjid, data = df) %>% mutate(geneStatus = geneStat) %>%
    mutate(visdat = as.POSIXct.Date(visdat, format = "%Y-%m-%d")) %>%
    arrange(visdat) %>%
    group_by(usubjid) %>%
    mutate(
      diffInDays = difftime(visdat, lag(visdat, 1, default = visdat[1]), unit = "days") %>%
        as.numeric() %>%
        cumsum()
    )
  
  df <- order_by(~usubjid, data = df)
  
  return(df)
}

addAgeColumn <- function(df, dat) {
  df <- df %>% mutate (
    age = age(enrollhd_demog$dm1__brthdtc[match(usubjid, enrollhd_demog$usubjid)], dat)
  )
  return(df)
}

addIscedColumn <- function(df) {
  df$cleanID <- gsub('["=]', '', df$usubjid)
  temp <- Baseline_Variable_Form_Data
  names(temp)[names(temp) == "usubjid"] <- "cleanID"
  temp <- temp %>% select(cleanID, isced)
  
  return(left_join(df, temp, by = "cleanID"))
}


addBaselineAgeColumn <- function(df) {
  df_first_visit <- df %>%
    group_by(usubjid) %>%
    slice(1) %>% select(age)
  
  df <- inner_join(df, df_first_visit, by = "usubjid") %>% rename(baselineAge = age.y)
  
  return(df)
}


mergeDataFrames <- function(isControl) {
  # df1 is baseline, df2 is follow ups
  if (isControl == 0) {
    df1 <- anti_join(BL_enrollhd_cognitive, enrollhd_cag, by = "usubjid") #all subjects that are not in cag
    df1 <- addAgeColumn(df1, df1$visdat)
    
    df2 <- anti_join(FU_enrollhd_cognitive, enrollhd_cag, by = "usubjid")
    df2 <- addAgeColumn(df2, df2$visdat)
    
    df <- setupDataframes(rbind(df1, df2), 0)
    df <- addIscedColumn(df)
    df <- addBaselineAgeColumn(df)
  }
  
  else {
    temp <- enrollhd_cag %>% select(usubjid)
    df1 <- inner_join(BL_enrollhd_cognitive, temp, by = "usubjid")
    df1 <- addAgeColumn(df1, df1$visdat)
    
    #adding age variable from demographic to follow up data frame
    df2 <- inner_join(FU_enrollhd_cognitive, temp, by = "usubjid")
    df2 <- addAgeColumn(df2, df2$visdat)
    
    df <- setupDataframes(rbind(df1, df2), 1)
    df <- addIscedColumn(df)
    df <- addBaselineAgeColumn(df)
  }
  return(df)
}

# set dataframes up
controlDf <- mergeDataFrames(0)
genePositiveDf <- mergeDataFrames(1)
fullDf <- rbind(genePositiveDf, controlDf)

# set factor levels
fullDf$geneStatus = as.factor(fullDf$geneStatus)
fullDf$geneStatus <- relevel(fullDf$geneStatus, ref = "0")
fullDf$isced <- as.factor(fullDf$isced)
fullDf$isced <- relevel(fullDf$isced, ref = "3")

# if AIC and BIC lower then that is better
sdmt1Model = lme(sdmt1 ~ diffInDays * geneStatus * isced * baselineAge, random = (~1|usubjid), data = fullDf, na.action = na.exclude)
ggpredict(sdmt1Model, c("diffInDays", "geneStatus", "isced", "baselineAge")) %>% plot()

sitModel = lme(sit1 ~ diffInDays * geneStatus * isced * baselineAge, random = (~1|usubjid), data = fullDf, na.action = na.exclude)
ggpredict(sitModel, c("diffInDays", "geneStatus", "isced", "baselineAge")) %>% plot()