---
title: "KNN_proximity_integration"
output: html_document
date: "2025-10-23"
---

## Setup

replace parameters below with the name of your metadata columns
Runs with a seurat object called seurat_obj

```{r}
# load integrated Mouse salamander TP73 neuron object 
seurat_obj<-readRDS("seurat5_clean_Allen_Woych_TP73_if5000_8000SVG_SCT_OG_d50_041325")

seurat_obj@meta.data$TP73_type <-factor(seurat_obj@meta.data$TP73_type, levels = c("pw-5","pw-0","pw-13","pw-1","pw-8","pw-6",
"pw-11","pw-7","pw-4","pw-2","pw-12","pw-3","pw-9","pw-10",
"ms-CR","ms-ETC","ms-PV"))
```






```{r}
###############################################################################
# 🧭 CONFIGURATION SECTION — EDIT THESE PARAMETERS ONLY
###############################################################################

# seurat_obj <-   # name of your integrated object
# important! celltype_col is one metadata column with the cell type annotations of the two species

celltype_col   <- "TP73_type"   # metadata column for cell types
condition_col  <- "species"              # metadata column for dataset/condition
condition_A    <- "mouse"                   # name of first condition (A)
condition_B    <- "pleuro"                  # name of second condition (B)
graph_name     <- "integrated_snn"          # Seurat graph name (e.g. "integrated_nn" or "SCT_nn")

output_prefix  <- paste0(condition_A, "_", condition_B, "_KNN")  # used for filenames
```


# Analysis

```{r}
###############################################################################
# 🧩 Load data and prepare metadata / graph
###############################################################################

library(Matrix)
library(dplyr)
library(tidyr)
library(pheatmap)

# Ensure seurat_obj exists
if (!exists("seurat_obj")) stop("❌ Please load your Seurat object as `seurat_obj` before running this script.")

# Extract metadata and graph
meta <- seurat_obj@meta.data
stopifnot(all(c(celltype_col, condition_col) %in% colnames(meta)))

if (!graph_name %in% names(seurat_obj@graphs))
  stop(paste("❌ Graph", graph_name, "not found in seurat_obj@graphs"))

knn_graph <- seurat_obj@graphs[[graph_name]]
knn_mat <- as(knn_graph, "dgCMatrix")

# Align metadata and graph
common <- intersect(rownames(knn_mat), rownames(meta))
meta <- meta[common, ]
knn_mat <- knn_mat[common, common]

# Clean up cell type names
meta[[celltype_col]] <- as.character(meta[[celltype_col]])
meta[[celltype_col]][is.na(meta[[celltype_col]]) | meta[[celltype_col]] == ""] <- "Unlabeled"

# Define condition-specific cells
cells_A <- rownames(meta)[meta[[condition_col]] == condition_A]
cells_B <- rownames(meta)[meta[[condition_col]] == condition_B]

message("✅ Cells loaded: ", length(cells_A), " in ", condition_A, " | ", length(cells_B), " in ", condition_B)

###############################################################################
# ⚙️ Function: compute directional proximity (A→B)
###############################################################################

get_directional_proximity <- function(source_cells, target_cells, meta, knn_mat, celltype_col) {
  proximity_list <- list()

  for (cell in source_cells) {
    if (!(cell %in% rownames(knn_mat))) next
    src_type <- meta[cell, celltype_col]

    nn_idx <- which(knn_mat[cell, ] > 0)
    if (length(nn_idx) == 0) next

    neighbors <- rownames(knn_mat)[nn_idx]
    neighbors_target <- neighbors[neighbors %in% target_cells]
    if (length(neighbors_target) == 0) next

    target_types <- meta[neighbors_target, celltype_col]
    counts <- table(target_types)

    df <- data.frame(
      source_type = src_type,
      target_type = names(counts),
      count = as.numeric(counts),
      stringsAsFactors = FALSE
    )
    proximity_list[[length(proximity_list) + 1]] <- df
  }

  if (length(proximity_list) == 0)
    stop("❌ No cross-dataset neighbors found in this direction.")

  proximity_df <- do.call(rbind, proximity_list)
  proximity_summary <- proximity_df %>%
    group_by(source_type, target_type) %>%
    summarise(count = sum(count), .groups = "drop")

  mat <- proximity_summary %>%
    tidyr::pivot_wider(names_from = target_type, values_from = count, values_fill = 0) %>%
    as.data.frame()

  rownames(mat) <- mat$source_type
  mat <- mat[, -1, drop = FALSE]

  # Normalize rows
  mat_norm <- sweep(mat, 1, rowSums(mat), "/")
  mat_norm[is.na(mat_norm)] <- 0
  return(mat_norm)
}

###############################################################################
# 🔁 Compute both directions (A→B and B→A)
###############################################################################

message("Computing ", condition_A, " → ", condition_B, " proximity...")
A_to_B <- get_directional_proximity(cells_A, cells_B, meta, knn_mat, celltype_col)

message("Computing ", condition_B, " → ", condition_A, " proximity...")
B_to_A <- get_directional_proximity(cells_B, cells_A, meta, knn_mat, celltype_col)

###############################################################################
# 🧮 Harmonize matrices and compute symmetric average
###############################################################################

all_A_types <- union(rownames(A_to_B), colnames(B_to_A))
all_B_types <- union(colnames(A_to_B), rownames(B_to_A))

pad_matrix <- function(mat, all_rows, all_cols) {
  missing_rows <- setdiff(all_rows, rownames(mat))
  if (length(missing_rows) > 0) {
    new_rows <- matrix(0, nrow = length(missing_rows), ncol = ncol(mat),
                       dimnames = list(missing_rows, colnames(mat)))
    mat <- rbind(mat, new_rows)
  }
  missing_cols <- setdiff(all_cols, colnames(mat))
  if (length(missing_cols) > 0) {
    new_cols <- matrix(0, nrow = nrow(mat), ncol = length(missing_cols),
                       dimnames = list(rownames(mat), missing_cols))
    mat <- cbind(mat, new_cols)
  }
  mat <- mat[all_rows, all_cols]
  mat[is.na(mat)] <- 0
  return(mat)
}

A_to_B <- pad_matrix(A_to_B, all_A_types, all_B_types)
B_to_A <- pad_matrix(B_to_A, all_B_types, all_A_types)

sym_matrix <- (A_to_B + t(B_to_A)) / 2
sym_matrix[is.na(sym_matrix)] <- 0
```



```{r}
sym_matrix<-sym_matrix[,c("pw-5","pw-0","pw-13","pw-1","pw-8","pw-6",
                 "pw-11","pw-7","pw-4","pw-2","pw-12","pw-3","pw-9","pw-10")]

sym_matrix<-t(sym_matrix)

```



```{r}
###############################################################################
# 🎨 Plot fixed-order symmetric similarity heatmap
###############################################################################

 common_order <- intersect(rownames(sym_matrix), colnames(sym_matrix))
 sym_matrix_plot <- sym_matrix[common_order, common_order, drop = FALSE]
 
pdf("TP73_Snn_Allen_pleuro_110925.pdf", width = 5, height = 10)
 
pheatmap(
  sym_matrix,
  cluster_rows = F,
  cluster_cols = F,
  main = paste("Symmetric", condition_A, "–", condition_B, "Cell Type Similarity (SNN Space)"),
  color = colorRampPalette(c("white", "#842B2B"))(100),
  border_color = NA,
  angle_col = 45,
  cellwidth = 25,
  cellheight = 25

)

dev.off()

```


```{r}
###############################################################################
# 💾 Save results
###############################################################################

write.csv(A_to_B,        paste0(output_prefix, "_AtoB.csv"), quote = FALSE)
write.csv(B_to_A,        paste0(output_prefix, "_BtoA.csv"), quote = FALSE)
write.csv(sym_matrix,    paste0(output_prefix, "_Symmetric.csv"), quote = FALSE)
write.csv(sym_matrix_plot, paste0(output_prefix, "_SymmetricOrdered.csv"), quote = FALSE)

message("✅ Done! Files saved with prefix: ", output_prefix)

```

