# install.packages(c("dodgr", "igraph", "ggplot2"))  # run once if needed

library(dodgr)
library(igraph)
library(ggplot2)
library(grid)  # for unit()

## ---------------------------------------------------------------
## 1. Build a perfect full binary tree of depth 6 + upstream X
## ---------------------------------------------------------------

depth <- 6

# Internal nodes (those that have children)
internal_indices <- 1:(2^(depth - 1) - 1)

edges_from <- character()
edges_to   <- character()

for (i in internal_indices) {
  edges_from <- c(edges_from, paste0("N", i), paste0("N", i))
  edges_to   <- c(edges_to,   paste0("N", 2 * i), paste0("N", 2 * i + 1))
}

# Upstream node X feeding N1
edges_from <- c("X", edges_from)
edges_to   <- c("N1", edges_to)

graph <- data.frame(
  edge_id  = seq_along(edges_from),
  from     = edges_from,
  to       = edges_to,
  weight   = 1,
  distance = 1,
  stringsAsFactors = FALSE
)

## ---------------------------------------------------------------
## 2. Define OD flows: from X to all leaves
## ---------------------------------------------------------------

root <- "X"
leaf_indices <- 2^(depth - 1):(2^depth - 1)
leaves <- paste0("N", leaf_indices)

from  <- root
to    <- leaves

n_leaves  <- length(leaves)
flows_vec <- c(rep(5,  n_leaves / 2),
               rep(10, n_leaves / 2))
flows <- matrix(flows_vec, nrow = 1)

## ---------------------------------------------------------------
## 3. Aggregate flows with dodgr
## ---------------------------------------------------------------

agg <- dodgr_flows_aggregate(
  graph     = graph,
  from      = from,
  to        = to,
  flows     = flows,
  contract  = FALSE,
  norm_sums = FALSE
)

## ---------------------------------------------------------------
## 4. Compute pressure drops: dP = R * Q^2 (R = 1 for demo)
## ---------------------------------------------------------------

agg$R_coeff    <- 1
agg$dp_edge_Pa <- agg$R_coeff * agg$flow^2

cat("Edge flows and dP:\n")
print(agg[, c("edge_id", "from", "to", "flow", "dp_edge_Pa")])

## ---------------------------------------------------------------
## 5. Build igraph and compute a tree layout
## ---------------------------------------------------------------

g <- graph_from_data_frame(
  d = agg[, c("from", "to", "flow", "dp_edge_Pa")],
  directed = TRUE
)

# Layout as a rooted tree with X at the top
coords <- layout_as_tree(
  g,
  root = which(V(g)$name == root),
  mode = "out"
)

nodes_df <- data.frame(
  name = V(g)$name,
  x    = coords[, 1],
  y    = -coords[, 2],  # flip y so root is at top
  stringsAsFactors = FALSE
)

## ---------------------------------------------------------------
## 6. Find index run: root -> leaf path with max total dP
## ---------------------------------------------------------------

# Identify leaves from the graph (nodes with no outgoing edges)
graph_leaves <- V(g)[degree(g, mode = "out") == 0]$name

all_paths <- unlist(
  lapply(graph_leaves, function(leaf) {
    all_simple_paths(g, from = root, to = leaf)
  }),
  recursive = FALSE
)

path_info <- lapply(all_paths, function(p) {
  nodes <- names(p)
  edge_pairs <- cbind(nodes[-length(nodes)], nodes[-1])

  dp_vals <- apply(edge_pairs, 1, function(e) {
    agg$dp_edge_Pa[agg$from == e[1] & agg$to == e[2]]
  })

  list(
    nodes     = nodes,
    edges     = edge_pairs,
    total_dP  = sum(dp_vals)
  )
})

total_dPs <- sapply(path_info, function(x) x$total_dP)
index_idx <- which.max(total_dPs)
index_run <- path_info[[index_idx]]

cat("\nIndex run (max total dP): ",
    paste(index_run$nodes, collapse = " -> "),
    "\nTotal dP: ", index_run$total_dP, " Pa\n", sep = "")

# Edge keys for index run (from->to)
index_edges <- data.frame(
  from = index_run$edges[, 1],
  to   = index_run$edges[, 2],
  stringsAsFactors = FALSE
)
index_keys <- paste(index_edges$from, index_edges$to, sep = "->")

## ---------------------------------------------------------------
## 7. Build edge dataframe with coordinates and labels
## ---------------------------------------------------------------

# Start from agg so flow and dP come through cleanly
edges_df_from <- merge(
  agg[, c("from", "to", "flow", "dp_edge_Pa")],
  nodes_df,
  by.x = "from",
  by.y = "name",
  all.x = TRUE
)
colnames(edges_df_from)[colnames(edges_df_from) %in% c("x", "y")] <- c("x_from", "y_from")

edges_df <- merge(
  edges_df_from,
  nodes_df,
  by.x = "to",
  by.y = "name",
  all.x = TRUE
)
colnames(edges_df)[colnames(edges_df) %in% c("x", "y")] <- c("x_to", "y_to")

# Ensure numeric
edges_df$flow       <- as.numeric(edges_df$flow)
edges_df$dp_edge_Pa <- as.numeric(edges_df$dp_edge_Pa)
edges_df$x_from     <- as.numeric(edges_df$x_from)
edges_df$y_from     <- as.numeric(edges_df$y_from)
edges_df$x_to       <- as.numeric(edges_df$x_to)
edges_df$y_to       <- as.numeric(edges_df$y_to)

# Midpoints for edge labels
edges_df$x_mid <- (edges_df$x_from + edges_df$x_to) / 2
edges_df$y_mid <- (edges_df$y_from + edges_df$y_to) / 2

edges_df$label <- paste0(
  "Q=", edges_df$flow,
  "\n",
  "dP=", round(edges_df$dp_edge_Pa, 3), " Pa"
)

# Mark index-run edges
edges_df$key          <- paste(edges_df$from, edges_df$to, sep = "->")
edges_df$is_index_run <- edges_df$key %in% index_keys

## ---------------------------------------------------------------
## 8. Plot with ggplot2, highlight index run in red, save PDF
## ---------------------------------------------------------------

p <- ggplot() +
  # non-index edges in grey
  geom_segment(
    data = subset(edges_df, !is_index_run),
    aes(x = x_from, y = y_from, xend = x_to, yend = y_to),
    arrow   = arrow(length = unit(2, "mm")),
    lineend = "round",
    colour  = "grey70",
    size    = 0.4
  ) +
  # index-run edges in red, thicker
  geom_segment(
    data = subset(edges_df, is_index_run),
    aes(x = x_from, y = y_from, xend = x_to, yend = y_to),
    arrow   = arrow(length = unit(3, "mm")),
    lineend = "round",
    colour  = "red",
    size    = 1.1
  ) +
  # edge labels (for all edges)
  geom_text(
    data = edges_df,
    aes(x = x_mid, y = y_mid, label = label),
    size = 2.0,
    lineheight = 0.9
  ) +
  # nodes
  geom_point(
    data = nodes_df,
    aes(x = x, y = y),
    size = 3
  ) +
  geom_text(
    data = nodes_df,
    aes(x = x, y = y, label = name),
    vjust = -1,
    size  = 3
  ) +
  theme_void() +
  ggtitle("Depth-6 Binary Tree with Flows, dP, and Index Run Highlighted")

ggsave("tree_flows_depth6_index_run.pdf", p, width = 8, height = 11)

cat("Saved PDF: tree_flows_depth6_index_run.pdf in ", getwd(), "\n")