options(scipen = 999)
setwd("C:/Users/irina/Documents/DND/EOHI/eohi2")

# Load required libraries
library(corrplot)
library(Hmisc)
library(psych)

# Load the data
exp2_data <- read.csv("eohi2.csv")

# Define the two sets of variables
set1_vars <- c("DGEN_past_5_mean", "DGEN_past_10_mean", "DGEN_fut_5_mean", "DGEN_fut_10_mean", 
               "DGEN_past_5.10_mean", "DGEN_fut_5.10_mean", "DGENpast_global_mean", 
               "DGENfut_global_mean", "DGEN_5_global_mean", "DGEN_10_global_mean", "DGEN_5.10_global_mean")
set2_vars <- c("aot_total", "crt_correct", "crt_int")

# Create subset with only the variables of interest
correlation_data <- exp2_data[, c(set1_vars, set2_vars)]

# ===== NORMALITY CHECKS =====
# Shapiro-Wilk tests for normality (n < 5000)
for(var in names(correlation_data)) {
  if(length(na.omit(correlation_data[[var]])) <= 5000) {
    shapiro_result <- shapiro.test(correlation_data[[var]])
    cat(sprintf("%s: Shapiro-Wilk p = %.5f %s\n", 
                var, shapiro_result$p.value, 
                ifelse(shapiro_result$p.value < 0.05, "(NOT normal)", "(normal)")))
  }
}


# Visual normality checks
pdf("normality_plots_domain_general_vars.pdf", width = 12, height = 8)
par(mfrow = c(2, 4))
for(var in names(correlation_data)) {
  # Histogram with normal curve overlay
  hist(correlation_data[[var]], main = paste("Histogram:", var), 
       xlab = var, freq = FALSE)
  curve(dnorm(x, mean = mean(correlation_data[[var]], na.rm = TRUE), 
              sd = sd(correlation_data[[var]], na.rm = TRUE)), 
        add = TRUE, col = "red", lwd = 2)
  
  # Q-Q plot
  qqnorm(correlation_data[[var]], main = paste("Q-Q Plot:", var))
  qqline(correlation_data[[var]], col = "red", lwd = 2)
}
dev.off()

# ===== LINEARITY CHECKS =====
# Check linearity between variable pairs
pdf("linearity_plots_domain_general_vars.pdf", width = 15, height = 10)
par(mfrow = c(4, 3))
for(i in 1:length(set1_vars)) {
  for(j in 1:length(set2_vars)) {
    var1 <- set1_vars[i]
    var2 <- set2_vars[j]
    
    # Scatter plot with regression line
    plot(correlation_data[[var1]], correlation_data[[var2]], 
         main = paste(var1, "vs", var2),
         xlab = var1, ylab = var2, pch = 21, cex = 0.6, bg = "lightblue", col = "black")
    
    # Add linear regression line
    lm_fit <- lm(correlation_data[[var2]] ~ correlation_data[[var1]])
    abline(lm_fit, col = "red", lwd = 2)
    
    # Add LOESS smooth line for non-linear pattern detection
    loess_fit <- loess(correlation_data[[var2]] ~ correlation_data[[var1]])
    x_seq <- seq(min(correlation_data[[var1]], na.rm = TRUE), 
                 max(correlation_data[[var1]], na.rm = TRUE), length = 100)
    loess_pred <- predict(loess_fit, x_seq)
    lines(x_seq, loess_pred, col = "blue", lwd = 2, lty = 2)
    
    # Calculate R-squared for linear fit
    r_squared <- summary(lm_fit)$r.squared
    cat(sprintf("%s vs %s: R² = %.5f\n", var1, var2, r_squared))
  }
}
dev.off()

# Residual analysis for linearity
pdf("residual_plots_domain_general_vars.pdf", width = 15, height = 10)
par(mfrow = c(4, 3))
for(i in 1:length(set1_vars)) {
  for(j in 1:length(set2_vars)) {
    var1 <- set1_vars[i]
    var2 <- set2_vars[j]
    
    lm_fit <- lm(correlation_data[[var2]] ~ correlation_data[[var1]])
    residuals <- residuals(lm_fit)
    fitted <- fitted(lm_fit)
    
    plot(fitted, residuals, 
         main = paste("Residuals:", var1, "vs", var2),
         xlab = "Fitted Values", ylab = "Residuals", pch = 21, cex = 0.6, bg = "lightblue", col = "black")
    abline(h = 0, col = "red", lwd = 2)
    
    # Add smooth line to residuals
    lines(lowess(fitted, residuals), col = "blue", lwd = 2)
  }
}
dev.off()


# Calculate Spearman correlation matrix only
cor_matrix_spearman <- cor(correlation_data, method = "spearman")

# Print correlation matrix with 5 decimal places
print(round(cor_matrix_spearman, 5))

# Separate correlations between the two sets (Spearman)
set1_set2_cor <- cor_matrix_spearman[set1_vars, set2_vars]
print(round(set1_set2_cor, 5))

# Calculate correlations within each set (Spearman)
set1_within_cor <- cor_matrix_spearman[set1_vars, set1_vars]
set2_within_cor <- cor_matrix_spearman[set2_vars, set2_vars]

# Statistical significance tests (Spearman)
cor_test_results_spearman <- rcorr(as.matrix(correlation_data), type = "spearman")

for(i in 1:length(set1_vars)) {
  for(j in 1:length(set2_vars)) {
    var1 <- set1_vars[i]
    var2 <- set2_vars[j]
    p_val <- cor_test_results_spearman$P[var1, var2]
    cat(sprintf("%s vs %s: p = %.5f\n", var1, var2, p_val))
  }
}

# Create correlation plot for Spearman only
pdf("correlation_plot_domain_general_vars_spearman.pdf", width = 10, height = 8)
corrplot(cor_matrix_spearman, method = "color", type = "upper", 
         order = "hclust", tl.cex = 0.8, tl.col = "black",
         addCoef.col = "black", number.cex = 0.7,
         title = "Spearman Correlation Matrix: Domain-General Vars vs Cognitive Measures")
dev.off()

# Summary statistics
desc_stats <- describe(correlation_data)
print(round(desc_stats, 5))

# Save results to CSV files
write.csv(round(cor_matrix_spearman, 5), "spearman_correlations_domain_general_vars.csv")
write.csv(round(desc_stats, 5), "descriptive_statistics_domain_general_vars.csv")

# Save correlation results in a formatted table
cor_results <- data.frame(
  Variable1 = character(),
  Variable2 = character(),
  Spearman_r = numeric(),
  P_value = numeric(),
  stringsAsFactors = FALSE
)

# Extract correlations between sets
for(i in 1:length(set1_vars)) {
  for(j in 1:length(set2_vars)) {
    var1 <- set1_vars[i]
    var2 <- set2_vars[j]
    r_val <- cor_matrix_spearman[var1, var2]
    p_val <- cor_test_results_spearman$P[var1, var2]
    
    cor_results <- rbind(cor_results, data.frame(
      Variable1 = var1,
      Variable2 = var2,
      Spearman_r = r_val,
      P_value = p_val,
      stringsAsFactors = FALSE
    ))
  }
}

write.csv(cor_results, "correlations - domain general vars.csv", row.names = FALSE)