# ============================================================
# CONDITIONAL TEST FOR NON-ADJACENT DEPENDENCY: A -> C | B
# For each observed trigram A_B_C, test whether:
#   P(C | A,B) > P(C | B)
# using one-sided Fisher's exact test
# ============================================================

library(tidyverse)

# ---------------------------
# 1. Read and reshape data
# ---------------------------
dat <- read.csv("input_data2.csv", check.names = FALSE)

names(dat) <- trimws(names(dat))
dat$Individual <- trimws(dat$Individual)

ng_cols <- grep("^n_gram", names(dat), value = TRUE)

long_dat <- dat %>%
  pivot_longer(
    cols = all_of(ng_cols),
    names_to = "slot",
    values_to = "ngram"
  ) %>%
  filter(!is.na(ngram), ngram != "") %>%
  mutate(slot_num = as.integer(stringr::str_extract(slot, "\\d+"))) %>%
  arrange(Individual, Model, slot_num)

# ---------------------------
# 2. Extract all contiguous trigrams
# ---------------------------
extract_trigrams <- function(seq_string) {
  units <- unlist(strsplit(seq_string, "_", fixed = TRUE))
  if (length(units) < 3) return(NULL)
  
  map_dfr(seq_len(length(units) - 2), function(i) {
    tibble(
      A = units[i],
      B = units[i + 1],
      C = units[i + 2],
      source_ngram = seq_string
    )
  })
}

tri_dat <- long_dat %>%
  mutate(trigrams = purrr::map(ngram, extract_trigrams)) %>%
  tidyr::unnest(trigrams)

# Optional:
# View total number of extracted trigrams
cat("Number of contiguous trigrams extracted:", nrow(tri_dat), "\n")

# ---------------------------
# 3. Conditional test function
# For a given middle unit B, and candidate A and C:
# compare:
#   P(C | A,B)
# vs
#   P(C | B and not A)
# which is equivalent to asking whether C depends on A given B
# ---------------------------
test_A_to_C_given_B <- function(tri_df, min_ab = 3, min_b = 10) {
  
  results <- list()
  
  B_values <- sort(unique(tri_df$B))
  
  idx <- 1
  
  for (b in B_values) {
    subB <- tri_df %>% filter(B == b)
    
    n_B <- nrow(subB)
    if (n_B < min_b) next
    
    A_values <- sort(unique(subB$A))
    C_values <- sort(unique(subB$C))
    
    for (a in A_values) {
      n_AB <- sum(subB$A == a)
      if (n_AB < min_ab) next
      
      for (c in C_values) {
        
        n_ABC <- sum(subB$A == a & subB$C == c)
        if (n_ABC == 0) next
        
        n_AB_notC      <- sum(subB$A == a & subB$C != c)
        n_notA_B_C     <- sum(subB$A != a & subB$C == c)
        n_notA_B_notC  <- sum(subB$A != a & subB$C != c)
        
        tab <- matrix(
          c(n_ABC, n_AB_notC,
            n_notA_B_C, n_notA_B_notC),
          nrow = 2,
          byrow = TRUE
        )
        
        # one-sided test: is C enriched after A_B?
        ft <- fisher.test(tab, alternative = "greater")
        
        p_C_given_AB <- n_ABC / n_AB
        p_C_given_B  <- (n_ABC + n_notA_B_C) / n_B
        delta        <- p_C_given_AB - p_C_given_B
        
        results[[idx]] <- tibble(
          A = a,
          B = b,
          C = c,
          n_B = n_B,
          n_AB = n_AB,
          n_ABC = n_ABC,
          P_C_given_AB = p_C_given_AB,
          P_C_given_B = p_C_given_B,
          delta = delta,
          odds_ratio = unname(ft$estimate),
          p_value = ft$p.value,
          table = paste0(
            "[[", n_ABC, ", ", n_AB_notC, "], [",
            n_notA_B_C, ", ", n_notA_B_notC, "]]"
          )
        )
        
        idx <- idx + 1
      }
    }
  }
  
  bind_rows(results) %>%
    mutate(
      p_adj_BH = p.adjust(p_value, method = "BH")
    ) %>%
    arrange(p_adj_BH, p_value, desc(delta))
}

cond_results <- test_A_to_C_given_B(tri_dat, min_ab = 3, min_b = 10)

# ---------------------------
# 4. Save full results
# ---------------------------
write.csv(cond_results, "conditional_nonadjacent_dependency_results.csv", row.names = FALSE)

# ---------------------------
# 5. Print strongest hits
# ---------------------------
top_hits <- cond_results %>%
  filter(delta > 0) %>%
  arrange(p_adj_BH, p_value, desc(delta))

print(top_hits, n = 30)

# ---------------------------
# 6. Minimal summary table for manuscript
# ---------------------------
summary_hits <- cond_results %>%
  filter(p_adj_BH < 0.05, delta > 0) %>%
  transmute(
    dependency = paste0(A, " -> ", C, " | ", B),
    n_B,
    n_AB,
    n_ABC,
    P_C_given_AB = round(P_C_given_AB, 3),
    P_C_given_B = round(P_C_given_B, 3),
    delta = round(delta, 3),
    p_adj_BH = signif(p_adj_BH, 3)
  )

print(summary_hits, n = 50)
write.csv(summary_hits, "conditional_nonadjacent_dependency_summary_hits.csv", row.names = FALSE)