esz_distribution_est.R

# estimating distributions

###########  Setup ########### 

# libraries
library(dplyr)
library(ggplot2)
library(ggrepel)
# library(ggridges)
library(lme4)
library(pwr) # for sample size calculations

# params 

estimate <- 'd'
base_dir <- '/Users/stephanienoble/Library/CloudStorage/GoogleDrive-s.noble@northeastern.edu/My Drive/Lab/Tasks-Ongoing/-K99/Effect_Size/'
data_dir <- paste0(base_dir, 'scripts/BrainEffeX_utils/inst/meta/v.RData')
# outlier_threshold <- 0.5

pooling_methods <- c('net','none')
save_plots <- TRUE

for (pooling_method in pooling_methods) {

combo_name <- paste0('pooling.', pooling_method, '.motion.regression.mv.none')
mv_combo_name <- paste0('pooling.', pooling_method, '.motion.none.mv.multi')

# setup

# load data
if (!exists("v")) {
  load(data_dir)
}

if (estimate == "d") {
  ci_lb <- "sim_ci_lb"
  ci_ub <- "sim_ci_ub"
} else if (estimate == "r_sq") {
  ci_lb <- "r_sq_sim_ci_lb"
  ci_ub <- "r_sq_sim_ci_ub"
}


# load data

# if v_orig doesn't exist, make v_orig <- v
if (!exists("v_orig")) {
  v_orig <- v
}
v <- v_orig # recreate v in case we messed anything up

# make output folder if doesn't exist

fn_basedir <- paste0(base_dir, '/manuscript/figures/plots/manuscript/curve_fit/',combo_name,'/')
if (!dir.exists(fn_basedir)) {
  dir.create(fn_basedir, recursive = TRUE)
}


###########  Summarize Data ########### 

# summarize studies

get_study_summaries <- function(data, study) {
  
  d_name <- numeric(length(data))
  d_mean <- numeric(length(data))
  d_sd <- numeric(length(data))
  d_cons_mean <- numeric(length(data))
  d_cons_sd <- numeric(length(data))
  d_n <- numeric(length(data))
  d_mv <- numeric(length(data))
  d_mv_lb <- numeric(length(data))
  d_mv_ub <- numeric(length(data))
  
  for (i in seq_along(data)) {
    
    # 1. Get point estimate mean and sd across vars
    
    # to get d, check whether dim is null (nested list)
    if (is.null(dim(data[[i]][[combo_name]][[estimate]]))) {
      d <- data[[i]][[combo_name]][[estimate]]
    } else {
      d <- data[[i]][[combo_name]][[estimate]][1,]
    }
    
    d_mean[i] <- mean(d)
    d_sd[i] <- sd(d)
    
    
    # 2. Get conservative estimate mean and sd across vars
    # from prep_data_for_plot.R
    # omitting na check in mean and sd and downsampling
    
    # unlist sim CIs if list
    if (is.list(data[[i]][[combo_name]][[ci_lb]])) {
      data[[i]][[combo_name]][[ci_lb]] <- unlist(data[[i]][[combo_name]][[ci_lb]])
      data[[i]][[combo_name]][[ci_ub]] <- unlist(data[[i]][[combo_name]][[ci_ub]])
    }
    
    na_idx <- is.na(data[[i]][[combo_name]][[estimate]]) | is.na(data[[i]][[combo_name]][[ci_lb]]) | is.na(data[[i]][[combo_name]][[ci_ub]])
    data[[i]][[combo_name]][[estimate]] <- data[[i]][[combo_name]][[estimate]][!na_idx]
    data[[i]][[combo_name]][[ci_lb]] <- data[[i]][[combo_name]][[ci_lb]][!na_idx]
    data[[i]][[combo_name]][[ci_ub]] <- data[[i]][[combo_name]][[ci_ub]][!na_idx]
    
    # sort data from smallest to largest effect size
    sorted_indices <- order(data[[i]][[combo_name]][[estimate]])
    sorted_estimate <- data[[i]][[combo_name]][[estimate]][sorted_indices]
    sorted_upper_bounds <- data[[i]][[combo_name]][[ci_ub]][sorted_indices]
    sorted_lower_bounds <- data[[i]][[combo_name]][[ci_lb]][sorted_indices]
    
    sorted_cons_estimate <- ifelse((abs(sorted_lower_bounds) > abs(sorted_upper_bounds)),
                                   ifelse((sorted_upper_bounds < 0),
                                          round(sorted_upper_bounds, 2), 0),
                                   ifelse((sorted_lower_bounds > 0),
                                          round(sorted_lower_bounds, 2), 0))
    
    
    d_cons_mean[i] <- mean(sorted_cons_estimate)
    d_cons_sd[i] <- sd(sorted_cons_estimate)
    
    
    # 3. Get sample size
    
    if (!is.null(data[[i]][[combo_name]]$n)) {
      d_n[i] <- data[[i]][[combo_name]]$n
    } else if (!is.null(data[[i]][[combo_name]]$n1) && !is.null(data[[i]][[combo_name]]$n2)) {
      d_n[i] <- data[[i]][[combo_name]]$n1 + data[[i]][[combo_name]]$n2
    } else {
      d_n[i] <- NA
    }
    
    # d_biased_sd[i] <- sd(d) * sqrt((length(d) - 1) / length(d)) # biased sd - doesn'd change results
    # TODO: catch normality test results
    # test for normality (usually light- or heavy-tailed, some approximately normal)
    # d_subset <- d[seq(1, length(d), length.out = 5000)] # max 5000 variables for shapiro test
    # d_normal_test[i] <- shapiro.test(d_subset)$p.value # p-value < 2e-16 -> not normal
    # qqnorm(d, pch=20, cex=0.5); qqline(d) # visualize
    
    # 4. Get multivariate effect size and bounds if available
    all_combos <- names(data[[i]])
    mv_combo_idx <- grep(mv_combo_name, all_combos)
    this_mv_combo <- all_combos[mv_combo_idx]
    d_mv[i] <- data[[i]][[this_mv_combo]]$d
    d_mv_lb[i] <- data[[i]][[this_mv_combo]]$sim_ci_lb
    d_mv_ub[i] <- data[[i]][[this_mv_combo]]$sim_ci_ub
  }
  
  # set up final data frame
  
  df <- data.frame(name = names(data), mean = d_mean, sd = d_sd, mean_cons = d_cons_mean, sd_cons = d_cons_sd, n = d_n, category = as.factor(study$category), dataset = I(study$dataset), ref = study$ref, mv = d_mv, mv_lb = d_mv_lb, mv_ub = d_mv_ub)
  
  # for meta: if exists, add n_studies
  if ("n_studies" %in% colnames(study)) {
    df$n_studies <- study$n_studies
  }
  
  # add reference type to dataset to distinguish data used for act from data used for FC (encompasses substantially different processing that should be nested within ref category)
  df <- df %>%
    mutate(dataset = paste(dataset, ref, sep = "_"))
  
  df <- df %>%
    mutate(overarching_category = case_when(
      category %in% c("biometric", "sex (demographic)", "age (demographic)") ~ "physical",
      category %in% c("cognitive", "psychiatric") ~ "psychological",
      category == "cognitive (task)" ~ "task (within-sub)",
      TRUE ~ "other"
    ))
  df$overarching_category <- as.factor(df$overarching_category)
  
  return(df)
  
}

# get summaries for orig
summary_data <- get_study_summaries(v$data, v$study)

# make separate data frame with conservative estimates (facilitates reuse of later functions on cons est)
summary_data_cons <- summary_data
summary_data_cons$mean <- summary_data_cons$mean_cons
summary_data_cons$sd <- summary_data_cons$sd_cons

# get summaries for meta

# first, for meta, category is inexplicably group_level - rename 
v$meta_category$study$category <- v$meta_category$study$group_level

# Assign all relevant datasets, n's, and study names to meta_category$study
v$meta_category$study$dataset <- vector("list", length(v$meta_category$study$name))
v$meta_category$study$each_n <- vector("list", length(v$meta_category$study$name))
v$meta_category$study$n <- vector("list", length(v$meta_category$study$name)) # number of unique subjects
v$meta_category$study$included_study_names <- vector("list", length(v$meta_category$study$name))
v$meta_category$study$n_studies <- integer(length(v$meta_category$study$name))

# For each meta-analysis, get all studies included, their associated datasets, and sample sizes
for (i in seq_along(v$meta_category$study$name)) {
  
  # get category and ref from meta_name
  meta_name <- v$meta_category$study$name[i]
  parts <- strsplit(meta_name, "_reference_")[[1]]
  meta_cat <- parts[1]
  meta_ref <- parts[2]

  matches <- which(v$study$category == meta_cat & v$study$ref == meta_ref)
  
  included_study_names <- v$study$name[matches]
  v$meta_category$study$included_study_names[[i]] <- included_study_names
  v$meta_category$study$dataset[[i]] <- v$study$dataset[matches]
  v$meta_category$study$overarching_category[[i]] <- unique(summary_data$overarching_category[matches])
  v$meta_category$study$each_n[[i]] <- summary_data$n[match(included_study_names, summary_data$name)]
  v$meta_category$study$n_studies[[i]] <- length(included_study_names)
  
  v$meta_category$study$n <- vector("list", length(v$meta_category$study$name))
  for (i in seq_along(v$meta_category$study$name)) {
    meta_datasets <- v$meta_category$study$dataset[[i]]
    unique_datasets <- unique(meta_datasets)
    unique_n <- 0
    for (ds in unique_datasets) {
      n_vals <- unlist(v$meta_category$study$each_n[[i]][v$meta_category$study$dataset[[i]] == ds])
      if (length(n_vals) > 0) {
        unique_n <- unique_n + max(n_vals, na.rm = TRUE)
      }
      # else do nothing (skip if no n)
    }
    
    # go through each field in v$meta_category$data
    for (field in names(v$meta_category$data[[i]])) {
      if (is.list(v$meta_category$data[[i]][[field]])) {
        # if list, subset to only those included in this meta-analysis
        v$meta_category$data[[i]][[field]]$n <- unique_n
      }
    }
    
  }
}
names(v$meta_category$study$dataset) <- v$meta_category$study$name
names(v$meta_category$study$each_n) <- v$meta_category$study$name
names(v$meta_category$study$included_study_names) <- v$meta_category$study$name
names(v$meta_category$study$n_studies) <- v$meta_category$study$name

# get summaries
summary_data__meta <- get_study_summaries(v$meta_category$data, v$meta_category$study)
summary_data_cons__meta <- summary_data__meta
summary_data_cons__meta$mean <- summary_data_cons__meta$mean_cons
summary_data_cons__meta$sd <- summary_data_cons__meta$sd_cons


# get total unique subjects across all studies
all_datasets <- unique(summary_data$dataset)
# remove hcp_voxel, which is a subset of hcp_shen_268
all_datasets <- all_datasets[all_datasets != "hcp_voxel"]
all_ns <- numeric(length(all_datasets))
for (j in seq_along(all_datasets)) {
  ds <- all_datasets[j]
  ns <- summary_data$n[summary_data$dataset == ds]
  if (length(ns) > 0) {
    all_ns[j] <- max(ns, na.rm = TRUE)
  } else {
    all_ns[j] <- 0
  }
}
total_n <- sum(all_ns, na.rm = TRUE)


########### Estimate Parameters & Plot Fits ########### 

# Function for plotting effect sizes (mean, sd, n) for each study - point est and conservative  
plot_study <- function(df, df_meta, main_title, fn, plot_type = "sd") {

print(paste0('Fitting lines for ', main_title))

# params
nmax_plt <- 10000000
nmax <- 1000000000
points_only_dir <- 'points_only/'
plot_fitted_lines <- TRUE

# Determine y variable and settings based on plot_type
if (plot_type == "sd") {
  y_var <- "sd"
  y_label <- "SD(d) across brain areas"
  fit_lines <- TRUE
  y_limits <- c(-0.05, 0.5)
} else if (plot_type == "mv") {
  y_var <- "mv"
  y_label <- "Multivariate effect size"
  # fit_lines <- FALSE
  y_limits <- NULL  # Let ggplot auto-scale
} else {
  stop("plot_type must be either 'sd' or 'mv'")
}

# sort by sample size
df <- df[order(df$n, decreasing = FALSE), ]

for (add_meta in c(TRUE, FALSE)) {

  if (add_meta) {
    meta_str <- '__meta'
    nmax_plt <- max(df_meta$n)
    alpha <- 0.15 # make non-meta points more transparent
  } else {
    meta_str <- ''
    nmax_plt <- max(df$n)
    alpha <- 0.7
  }
  
  # Only fit lines for sd plots
  # fit_options <- if (fit_lines) c(FALSE, TRUE) else c(FALSE)
  
  # for (plot_fitted_lines in fit_options) {
    
    # if (plot_fitted_lines && fit_lines) {
    if (plot_fitted_lines) {
      fitlines_str <- '__fitlinesY'
      
      if (plot_type == "mv") {
        
        # set up n
        n_seq <- data.frame(n=seq(0, nmax_plt, length.out=1000))
        n_seq[length(n_seq$n)+1,] <- nmax # ensure max n is included
        
        # run for each overarching category
        unique_cats <- unique(df$overarching_category)
        predicted_mv_cat <- vector("list", length(unique_cats))
        names(predicted_mv_cat) <- unique_cats
        
        # preallocate beta
        beta <- vector("list", length(unique_cats))
        max_mv <- vector("list", length(unique_cats))
        res <- vector("list", length(unique_cats))
        names(beta) <- unique_cats
        
        for (cat in unique_cats) {
          
          df_cat <- df[df$overarching_category == cat & !is.na(df$mv), ]
          
          n_obs <- nrow(df_cat)
          n_levels <- length(unique(df_cat$dataset))
          
          # if only one unique dataset, do without random effect
          if (length(unique(df_cat$dataset)) > 1) {
            
            # fit
            fit_cat <- lmer(mv ~ 1 + (1|dataset), data = df_cat)
            
            # Use fixed effects for prediction
            d <- as.matrix(model.matrix(~ 1, data = n_seq))
            beta[[cat]] <- fixef(fit_cat)
            preds <- as.vector(d %*% beta[[cat]])
            se <- sqrt(diag(d %*% vcov(fit_cat) %*% t(d)))
            lwr <- preds - 1.96 * se
            upr <- preds + 1.96 * se
            
            predicted_mv_cat[[cat]] <- cbind(fit = preds, lwr = lwr, upr = upr)
            
          } else {
            fit_cat <- lm(mv ~ I(1/sqrt(n)), data=df_cat)
            predicted_mv_cat[[cat]] <- predict(fit_cat, n_seq, interval="confidence")
            beta[[cat]] <- coef(fit_cat)
          }
          
          max_mv[[cat]] <- predicted_mv_cat[[cat]][length(n_seq$n),]
          
          res[[cat]] <- c(beta[[cat]], max_mv[[cat]])
          
        }
        
        
      } else {
        
        # # fit curves to each category
        # fit <- lm(sd ~ I(1/sqrt(n)), data=df)
        # predicted_sd <- predict(fit, df)
        # ids_above <- df$sd > predicted_sd
        # fit_high <- lm(sd ~ I(1/sqrt(n)), data=df[ids_above,])
        # fit_low <- lm(sd ~ I(1/sqrt(n)), data=df[!ids_above,])
        # 
        # # interpolate predictions for all n
        n_seq <- data.frame(n=seq(0, nmax_plt, length.out=1000))
        n_seq[length(n_seq$n)+1,] <- nmax # ensure max n is included
        # predicted_sd_seq <- predict(fit, n_seq, interval="confidence")
        
        # run for each overarching category
        unique_cats <- unique(df$overarching_category)
        predicted_sd_cat <- vector("list", length(unique_cats))
        names(predicted_sd_cat) <- unique_cats
        
        # preallocate beta
        beta <- vector("list", length(unique_cats))
        max_sd <- vector("list", length(unique_cats))
        res <- vector("list", length(unique_cats))
        names(beta) <- unique_cats
        
        for (cat in unique_cats) {
          
          df_cat <- df[df$overarching_category == cat, ]
          
          n_obs <- nrow(df_cat)
          n_levels <- length(unique(df_cat$dataset))
          
          # if only one unique dataset, do without random effect
          if (length(unique(df_cat$dataset)) > 1) {
            
            # fit_cat <- lmer(sd ~ I(1/sqrt(n)) + (1|dataset), data=df_cat)
            
            # Use only fixed effects for prediction
            fit_cat <- lmer(sd ~ I(1/sqrt(n)) + (1|dataset), data = df_cat)
            d <- model.matrix(~ I(1/sqrt(n)), data = n_seq)
            beta[[cat]] <- fixef(fit_cat)
            preds <- as.vector(d %*% beta[[cat]])
            se <- sqrt(diag(d %*% vcov(fit_cat) %*% t(d)))
            lwr <- preds - 1.96 * se
            upr <- preds + 1.96 * se
            
            predicted_sd_cat[[cat]] <- cbind(fit = preds, lwr = lwr, upr = upr)
            
          } else {
            fit_cat <- lm(sd ~ I(1/sqrt(n)), data=df_cat)
            predicted_sd_cat[[cat]] <- predict(fit_cat, n_seq, interval="confidence")
            beta[[cat]] <- coef(fit_cat)
          }
          
          max_sd[[cat]] <- predicted_sd_cat[[cat]][length(n_seq$n),]
          
          res[[cat]] <- c(beta[[cat]], max_sd[[cat]])
          
        }
      }
      
      # remove null entries from res and convert to data frame
      res <- data.frame(res[!sapply(res, is.null)])

      fn_plot <- fn
      
    } else {
      # If not plotting fitlines, change output folder to 'points_only'
      fn_plot <- file.path(dirname(fn), points_only_dir, basename(fn))
      if (!dir.exists(dirname(fn_plot))) {
        dir.create(dirname(fn_plot), recursive = TRUE)
      }
      if (plot_type == "mv") {
        res <- NULL
      } else {
        res <- NULL
      }
    }
    
    # plot
    # ggplot2 implementation
    df$sqrt_n <- sqrt(df$n)
    df_meta$sqrt_n <- sqrt(df_meta$n)
    cats <- levels(df$overarching_category)
    color_map <- setNames(RColorBrewer::brewer.pal(length(cats), "Set1"), cats)

    # Create base plot with dynamic y variable
    p <- ggplot(df, aes_string(x = "sqrt_n", y = y_var, color = "overarching_category")) +
      geom_point(size = 1.5, alpha = alpha, stroke = 0) +
      scale_color_manual(values = color_map) +
      labs(title = main_title,
            x = "Sqrt(n)",
            y = y_label,
            color = "Category") +
      theme_bw() +
      theme(
        plot.title = element_text(hjust = 0.5),
        legend.position = c(0.98, 0.98),
        legend.justification = c("right", "top"),
        legend.key.size = unit(0.7, "lines"),
        panel.grid.major = element_blank(),
        panel.grid.minor = element_blank()
      ) +
      scale_x_continuous(limits = c(0, max(sqrt(nmax_plt), max(sqrt(df_meta$n), na.rm=TRUE), max(sqrt(df$n), na.rm=TRUE))), expand = c(0, 0)) +
      geom_hline(yintercept = 0, linetype = "dashed", color = "black", linewidth = 0.3)
    
    # Add y-axis limits only for sd plots
    if (!is.null(y_limits)) {
      p <- p + scale_y_continuous(limits = y_limits, expand = c(0, 0))
    }
    
    # Add mean points only for sd plots
    if (plot_type == "sd") {
      p <- p + geom_point(aes(x = sqrt_n, y = mean), size = 1.5, alpha = alpha, stroke = 0)
    }

    # Add fitted lines only for sd plots when requested
    # if (plot_fitted_lines && fit_lines) {
    if (plot_fitted_lines) {
      
      if (plot_type == "mv") {
        for (cat in unique_cats) {
          pred_mat <- predicted_mv_cat[[cat]]
          if (!is.null(pred_mat) && is.matrix(pred_mat) && all(c("fit","lwr","upr") %in% colnames(pred_mat)) && nrow(pred_mat) == length(n_seq$n)) {
            print(unique(pred_mat))
            pred_df <- data.frame(sqrt_n = sqrt(n_seq$n),
                                  fit = pred_mat[,"fit"],
                                  lwr = pred_mat[,"lwr"],
                                  upr = pred_mat[,"upr"],
                                  overarching_category = cat)
            p <- p +
              geom_line(data = pred_df, aes(x = sqrt_n, y = fit, color = overarching_category), linewidth = 1, alpha = alpha) +
              geom_line(data = pred_df, aes(x = sqrt_n, y = lwr, color = overarching_category), linetype = "dotted", linewidth = 0.8, alpha = alpha) +
              geom_line(data = pred_df, aes(x = sqrt_n, y = upr, color = overarching_category), linetype = "dotted", linewidth = 0.8, alpha = alpha)
          }
        }
        
      } else {
        for (cat in unique_cats) {
          pred_mat <- predicted_sd_cat[[cat]]
          if (!is.null(pred_mat) && is.matrix(pred_mat) && all(c("fit","lwr","upr") %in% colnames(pred_mat)) && nrow(pred_mat) == length(n_seq$n)) {
            pred_df <- data.frame(sqrt_n = sqrt(n_seq$n),
                                  fit = pred_mat[,"fit"],
                                  lwr = pred_mat[,"lwr"],
                                  upr = pred_mat[,"upr"],
                                  overarching_category = cat)
            p <- p +
              geom_line(data = pred_df, aes(x = sqrt_n, y = fit, color = overarching_category), linewidth = 1, alpha = alpha) +
              geom_line(data = pred_df, aes(x = sqrt_n, y = lwr, color = overarching_category), linetype = "dotted", linewidth = 0.8, alpha = alpha) +
              geom_line(data = pred_df, aes(x = sqrt_n, y = upr, color = overarching_category), linetype = "dotted", linewidth = 0.8, alpha = alpha)
          }
        }
      }
    }

    if (add_meta && nrow(df_meta) > 0) {
      # Use gsub to format label with line break and study count
      df_meta$label <- paste0(gsub("_reference_", " (", df_meta$name), ")\n", df_meta$n_studies, " ", ifelse(df_meta$n_studies == 1, "study", "studies"))
      p <- p +
        geom_point(data = df_meta, aes_string(x = "sqrt_n", y = y_var, color = "overarching_category"), shape = 8, size = 2) +
        geom_text_repel(data = df_meta, aes_string(x = "sqrt_n", y = y_var, label = "label"), size = 2.2, segment.size = 0.3, force = 10, max.overlaps = Inf)
    }

    # Create filename suffix based on plot type and fitted lines
    type_suffix <- if (plot_type == "mv") "__mv" else ""
    # fitlines_suffix <- if (plot_fitted_lines && fit_lines) "__fitlinesY" else ""
    fitlines_suffix <- if (plot_fitted_lines) "__fitlinesY" else ""
    this_fn <- paste0(fn_plot, type_suffix, fitlines_suffix, meta_str, '.pdf')
    if (save_plots) {
      ggsave(this_fn, plot = p, width = 5, height = 4)
    } else {
      show(p)
    }
    
  # }
}

return(res)
}

plot_non_mv <- FALSE # TODO - tmp

if (plot_non_mv) {
# get point est and conservative est
res <- plot_study(summary_data, summary_data__meta, "Point Estimates", paste0(fn_basedir,'point_est'))
res_cons <- plot_study(summary_data_cons, summary_data_cons__meta, "Conservative Estimates", paste0(fn_basedir,'cons_est'))
  
# repeat, but for large N studies
res_large <- plot_study(summary_data[summary_data$n > 900,], summary_data__meta[summary_data__meta$n > 900,],"Point Estimates (n > 900)", paste0(fn_basedir,'point_estimates_n900'))
}

# get multivariate effect sizes
res_mv <- plot_study(summary_data, summary_data__meta, "Multivariate Effect Sizes", paste0(fn_basedir,'mv_est'), plot_type = "mv")


#     outliers <- (df_orig$mean > mean(df_orig$mean) + outlier_threshold*sd(df_orig$mean)) | (df_orig$mean < mean(df_orig$mean) - outlier_threshold*sd(df_orig$mean))
#     df <- df_orig[!outliers,]




###########  Projected Density Plots ########### 

# make effect size distribution plots for each overarching category

# Overlay density plots for all categories, using default factor colors
#cats <- levels(summary_data$overarching_category)
cats <- c("physical", "psychological", "task..within.sub.")
# Use default R colors for factor levels
cat_factor <- factor(cats)
cat_colors <- as.numeric(cat_factor)

fn <- paste0(fn_basedir, 'projected_distributions.pdf')
if (save_plots) {
pdf(file=fn, width=5, height=4)
}

d <- seq(-1, 1, length.out=1000)
plot(d, rep(0, length(d)), type = "n",
     xlab = "d", ylab = "Density", main = "Density Plot by Category",
     xlim = c(-1,1), ylim = c(0, 50))

sigmas_master <- NULL


## Overlapping

for (i in seq_along(cats)) {
  
  cat <- cats[i]
  color <- cat_colors[i]
  
  sigmas <- c(res_cons["lwr",cat],res_cons["(Intercept)",cat],res["(Intercept)",cat],res["upr",cat])
  names(sigmas) <- c('lb_lb','lb','ub','ub_ub')
  
  for (name in names(sigmas)) {
    s <- sigmas[name]
    lty <- if (name %in% c("lb_lb", "ub_ub")) 2 else 1
    lwd <- if (name %in% c("lb_lb", "ub_ub")) 0.25 else 1
    if (s >= 0) {
      y <- dnorm(d, mean = 0, sd = s)
      lines(d, y, col = color, lwd = lwd, lty = lty)
    } else {
      segments(x0 = 0, y0 = 0, x1 = 0, y1 = 100, col = color, lwd = lwd, lty = lty)
    }
  }
  sigmas_master <- rbind(sigmas_master,sigmas)
}

legend('topright', legend=cats, col=cat_colors, lwd=2, cex=0.8, bty="n")
rownames(sigmas_master) <- cats
if (save_plots) {
  dev.off()
}


# save sigmas_master to file
if (save_plots) {
  write.csv(sigmas_master, file=paste0(fn_basedir, 'projected_distributions__sigmas.csv'), row.names=TRUE)
}

## Separate panels per category

panel_colors <- cat_colors
panel_names <- cats

fn <- paste0(fn_basedir, 'projected_distributions__panels.pdf')
if (save_plots) {
pdf(fn, width = 12, height = 4)  # width and height in inches
}

d <- seq(-1, 1, length.out=1000)
par(mfrow = c(1,length(panel_names))) # one row per category
for (i in seq_along(panel_names)) {
  cat <- panel_names[i]
  color <- panel_colors[i]
  sigmas <- sigmas_master[cat, ]
  # get max_y only from 'lb' curve for this category
  s_lb <- sigmas['lb']
  if (!is.na(s_lb) && s_lb >= 0) {
    y_lb <- dnorm(d, mean = 0, sd = s_lb)
    max_y <- max(y_lb, na.rm=TRUE)
  } else {
    max_y <- 1
  }
  plot(d, rep(0, length(d)), type = "n",
       xlab = "d", ylab = "Density", main = paste("Density Curves -", cat),
       xlim = c(-1,1), ylim = c(0, max_y * 1.05))
  for (j in seq_along(sigmas)) {
    s <- sigmas[j]
    lty <- if (names(sigmas)[j] %in% c("lb_lb", "ub_ub")) 2 else 1
    lwd <- if (names(sigmas)[j] %in% c("lb", "ub")) 1.5 else 1.5
    if (s >= 0) {
      y <- dnorm(d, mean = 0, sd = s)
      lines(d, y, col = color, lwd = lwd, lty = lty)
    } else {
      segments(x0 = 0, y0 = 0, x1 = 0, y1 = 100, col = color, lwd = lwd, lty = lty)
    }
  }
  legend('topright', legend=names(sigmas), col=rep(color, length(sigmas)), lwd=c(2,1,2,1), lty=c(1,2,1,2), cex=0.8, bty="n")
}
par(mfrow = c(1,1)) # reset layout

if (save_plots) {
dev.off()
}




##### REQUIRED N PLOTS #####
  
# for each category, plot required n to detect effects of different sizes with 80% power

alpha <- 0.05/2 # TODO: confirm
# alpha <- 0.05/2/35778 # corrected

# for saving plots
fn <- paste0(fn_basedir, 'required_n_distributions__panels.pdf') # name of output file
if (save_plots) {
pdf(fn, width = 5, height = 15)  # width and height in inches
}

cats <- levels(summary_data$overarching_category)
# Use default R colors for factor levels
cat_factor <- factor(cats)
cat_colors <- as.numeric(cat_factor)

# load sigmas_master again
sigmas_master = read.csv(paste0(fn_basedir, 'projected_distributions__sigmas.csv'), row.names = 1)
panel_colors <- cat_colors
panel_names <- cats

d <- seq(-1, 1, length.out=1000) # effect sizes to consider, same as above

# calculate minimum n needed to detect each d with 80% power, for one-sample and two-sample t-tests
# n_detect_one is for one-sample, n_detect_two is for two-sample
# initialize empty numeric vectors
n_detect_one <- numeric(length(d))
n_detect_two <- numeric(length(d))

# loop through d values and calculate n for each d value for one-sample and two-sample options
for (i in seq_along(d)) {
  n_detect_two[i] <- pwr.t.test(power = 0.8, d = d[i], sig.level = alpha, type = "two.sample")$n
  n_detect_one[i] <- pwr.t.test(power = 0.8, d = d[i], sig.level = alpha, type = "one.sample")$n
}

# setup for plotting multiple plots in one figure
par(mfrow = c(length(panel_names),1)) # one row per category

# loop through categories and plot
for (i in seq_along(panel_names)) {
  cat <- panel_names[i] # category name

  # if category name is "task..within.sub.", use n_detect_one, else n_detect_two
  if (cat == "task (within-sub)") {
    n_option <- n_detect_one
  } else {
    n_option <- n_detect_two
  }

  # set colors as above and get sigmas for that category
  color <- panel_colors[i]
  sigmas <- sigmas_master[cat, ]
    
  # get y values for upper bound sigma
  if (cat == "task (within-sub)") {
    cat = "task..within.sub."
  }
  y <- dnorm(d, mean = 0, sd = sigmas_master[cat,"ub"])

  sqrt_n_option <- sqrt(n_option)
  
  do_bin <- TRUE
  if (do_bin) {
    
    n_bins <- c(0, 25, 50, 100, 500, 1000, 5000, 50000, Inf)
    bin_labels <- paste(head(n_bins, -1), n_bins[-1], sep = "–")
    bin_indices <- cut(n_option, breaks = n_bins, include.lowest = TRUE, labels = bin_labels)
    
    # Sum y values in each bin
    binned_sums <- tapply(y, bin_indices, sum, na.rm = TRUE)
    # make binned_sums a density (AUC sums to 1)
    binned_sums <- binned_sums / sum(binned_sums, na.rm = TRUE)
    # make cumulative, skipping entries that are na
    binned_sums[!is.na(binned_sums)] <- cumsum(binned_sums[!is.na(binned_sums)])
    
    # Plot
    plot(seq_along(bin_labels), binned_sums, type = "n",
         xaxt = "n",
         xlab = "Minimum sample size to detect effects with 80% power",
         ylab = "Proportion of Effects", main = paste("Density Curves -", cat, " (required n for 80% power)")
    )
    axis(1, at = seq_along(bin_labels), labels = bin_labels)
    points(seq_along(bin_labels), binned_sums, col = color, pch = 19)
    lines(seq_along(bin_labels), binned_sums, col = color, lwd = 2)
    legend('topright', legend=names(sigmas), col=rep(color, length(sigmas)), lwd=c(2,1,2,1), lty=c(1,2,1,2), cex=0.8, bty="n")
  } else {
    # make binned_sums a density (AUC sums to 1)
    y_dens <- y / sum(y, na.rm = TRUE)
    # make cumulative, skipping entries that are na
    y_dens[!is.na(y_dens)] <- cumsum(y_dens[!is.na(y_dens)])
    # plot required n's vs. density
    plot(sqrt_n_option, y_dens, type = "n",
         xlab = "Minimum sample size to detect effects with 80% power", 
         ylab = "Density", main = paste("Density Curves -", cat, " (required n for 80% power)")
         )
    # add lines
    lines(sqrt_n_option, y_dens, col = color, lwd = 2)
  }
  
}

par(mfrow = c(1,1)) # reset layout

if (save_plots) {
dev.off()
}

}





# ADDL NOTES

# what is the estimated distribution of null cohen's d effect sizes at a given sample size?
# simulate null distributions of cohen's d for different sample sizes

# expected variance of estimated d, from d variance:
# estimate d_var for vector of n's
# n_vec <- summary_data$n
# d_var_vec <- (n_vec - 1) / (n_vec * (n_vec - 3))
# # d_var_vec2 <- (2*n_vec) / (n_vec - 2) + (d_mean^2) / (2*(n_vec - 1)) # from Hedges & Olkin (1985)
# d_var_r_vec <- 4 / (n_vec - 3) # from Borenstein et al. (2009), eq 6.7
# lines(sqrt(d_var_r_vec), type='b', col='purple', pch=20) # add to previous plot

# d_var <- function(n) {
#   return( (2*n) / (n - 2) + (d_mean^2) / (2*(n - 1)) ) # from Hedges & Olkin (1985)
# }


# re normality: for reference, here's a qqnorm for data randomly sampled from normal
# set.seed(42)
# t3 <- rnorm(5000, mean=mean(d), sd=sd(d))
# qqnorm(t3, pch=20); qqline(t3)