##########################################################################
### Example Code for the LLRT Method
##########################################################################

require(pacman)
pacman::p_load(data.table, equate, psych, dplyr, magrittr)

## FUNCTION INPUT
## FULL DATA - DATA.TABLE INCLUDING DATA FOR BOTH STUDIES
## COLUMNS NEEDED:
##   - "study", REFERENCE STUDY SHOULD BE STUDY == 1
##   - "frelated", HARMONIZED RELATED LATENT TRAIT SCORES
##   - "sfsubdomain", UNHARMONIZED SUBDOMAIN TRAIT SCORES

## FUNCTION OUTPUT
## DATASET WITH COLUMN ADDED:
##   - "subdomian_e_f", LLRT HARMONIZED SUBDOMAIN SCORES

llrt_singlestrata <- function(data){
  
  message('Data Prep')
  ## LINEAR MODEL
  total_dt <- copy(data)
  total_dt <- total_dt[!is.na(frelated)]
  
  # DATA PROCESSING FOR EQUATING --------------------------------------------
  
  message("Equating")
  total_dt[, `:=` (related = round(frelated*10, 0), subdomain = round(sfsubdomain*10, 0))]
  #total_dt[, `:=` (related = related - min(related), subdomain = subdomain - min(subdomain))]
  
  all_scores <- min(min(total_dt[,related], na.rm = T), min(total_dt[,subdomain], na.rm = T)):max(max(total_dt[,related], na.rm = T), max(total_dt[,subdomain], na.rm = T))
  template_dt <- data.table(id = all_scores)
  
  related_table <- total_dt %>%
    count(related) %>%
    filter(!is.na(related)) %>%
    rename(id = related, related_val = n)
  subdomain_table <- total_dt %>%
    count(subdomain) %>%
    filter(!is.na(subdomain)) %>%
    rename(id = subdomain, subdomain_val = n)
  
  total_table <- Reduce(function(x,y) merge(x,y,by="id",all=T), list(template_dt, related_table, subdomain_table))
  
  # DO LINEAR EQUATING ------------------------------------------------------
  
  equated <- equate(x = as.freqtab(total_table[, .(id, subdomain_val)]),
                    y = as.freqtab(total_table[, .(id, related_val)]),
                    type = "l", bootse = F, reps = 50)
  total_table[, subdomain_equated := equated$concordance$yx]
  
  # MERGE SCORES ------------------------------------------------------------
  
  post_equated <- merge(total_dt, total_table[, .(subdomain = id, subdomain_equated)], by = "subdomain", all.x = T)
  post_equated <- post_equated[!is.na(subdomain) & !is.na(related)]
  
  
  return(post_equated)
}

llrt <- function(full_data){
  
  full_data[, reference := as.numeric(study == 1)]
  reference_eq <- llrt_singlestrata(full_data[reference == 1])
  other_eq <- llrt_singlestrata(full_data[reference == 0])
  post_equated <- rbind(reference_eq, other_eq)
  
  # TRANSFORM SCORES ---------------------------------------------------------
  
  ## TRANSFORM SO THAT HAS MEAN 50 SD 10 IN REFERENCE SAMPLE
  
  message("Transforming Scores")
  ## Get to mean zero
  mean_adj <- post_equated[reference == 1, mean(subdomain_equated)]
  post_equated[, subdomain_e_f := subdomain_equated-mean_adj]
  
  sd_adj <- 10/post_equated[reference == 1, sd(subdomain_equated)]
  post_equated[, subdomain_e_f := subdomain_e_f*sd_adj + 50]
  
  return(post_equated)
}

## EXAMPLE OF HOW TO CALL FUNCTION
example_result <- llrt(full_data)