# Libraries
library(NlsyLinks)
library(discord)
library(utils)
library(tidyverse)
library(ggplot2)
# Set random seed for reproducibility
set.seed(1492)

# Disable scientific notation for clarity
options(scipen = 999)



conditions <- expand.grid(
  total_pairs = c(100, 250, 500, 750, 1000),
  relatedness = c(1, .5),
  cov_a = c(0, 0.25),
  cov_c = c(0, 0.25),
  cov_e = c(0, 0.25),
  ace_a = c(1),
  ace_c = c(1),
  ace_e = c(1)
)

set.seed(1492) # Set seed for reproducibility
n_trials <- 100
FAST <- FALSE # Set to FALSE for slower, more detailed analysis
results_list <- list()
name.results <- c("coef_xdiff", "p_xdiff", "r.squared")

for (cond in seq_along(conditions)) {
  current <- conditions[cond, ]
  temp_results <- matrix(NA, nrow = n_trials, ncol = length(name.results))
  colnames(temp_results) <- name.results

  for (i in 1:n_trials) {
    trial <- kinsim(
      r_vector = rep(current$relatedness, each = current$total_pairs),
      npg_all = current$total_pairs,
      ace_all = c(current$ace_a, current$ace_c, current$ace_e),
      cov_a = current$cov_a,
      cov_c = current$cov_c,
      cov_e = current$cov_e,
      variables = 2
    )

    extract <- data.frame(
      id = trial$id, r = trial$r,
      y_s1 = trial$y1_1, y_s2 = trial$y1_2,
      x_s1 = trial$y2_1, x_s2 = trial$y2_2
    )

    if (FAST == TRUE) {
      # faster
      # double enter the data
      extract2 <- rbind(
        transform(extract,
          y_s1 = y_s2, y_s2 = y_s1,
          x_s1 = x_s2, x_s2 = x_s1
        ),
        extract
      )
      extract2$y_diff <- extract2$y_s1 - extract2$y_s2
      extract2$x_diff <- extract2$x_s1 - extract2$x_s2
      extract2$x_bar <- (extract2$x_s1 + extract2$x_s2) / 2
      extract2$y_bar <- (extract2$y_s1 + extract2$y_s2) / 2

      # select pair with ydiff > 0
      extract3 <- extract2[extract2$y_diff > 0, ]

      fit <- tryCatch(
        lm(y_diff ~ x_bar + y_bar + x_diff, data = extract3),
        error = function(e) {
          return(NULL)
        }
      )
    }
    # slower
    if (FAST == FALSE) {
      fit <- tryCatch(
        discord_regression(
          data = extract, outcome = "y", predictors = "x",
          id = "id",
          sex = NULL,
          race = NULL,
          fast = TRUE
        ),
        error = function(e) {
          return(NULL)
        }
      )
    }
    if (!is.null(fit)) {
      sm <- summary(fit)
      temp_results[i, "coef_xdiff"] <- coef(sm)["x_diff", "Estimate"]
      temp_results[i, "p_xdiff"] <- coef(sm)["x_diff", "Pr(>|t|)"]
      temp_results[i, "r.squared"] <- sm$r.squared
    }
  }

  results_list[[cond]] <- as.data.frame(temp_results)
}

power_summary <- lapply(results_list, function(res) {
  data.frame(
    power_xdiff = mean(res$p_xdiff < 0.05, na.rm = TRUE),
    median_r2   = median(res$r.squared, na.rm = TRUE)
  )
})

final_results <- cbind(conditions, do.call(rbind, power_summary))