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

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

## ---------------------------------------------------------------
## 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: random per-leaf flows with large variation
## ---------------------------------------------------------------

set.seed(123)   # set or remove for reproducible randomness

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

n_leaves <- length(leaves)

# RANDOM LEAF FLOW RATES (choose range as you like)
min_flow <- 3
max_flow <- 50

flows_vec <- runif(n_leaves, min = min_flow, max = max_flow)

# dodgr requires flows as a matrix: nrow = length(from) = 1
flows <- matrix(flows_vec, nrow = 1)

cat("\nRandom per-leaf flows (units e.g. L/s or m³/s):\n")
print(round(flows_vec, 3))

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

agg <- dodgr_flows_aggregate(
  graph     = graph,
  from      = root,
  to        = leaves,
  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("\nEdge 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 & per-path total dP
## ---------------------------------------------------------------

# Leaves in 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=", round(edges_df$flow, 3),
  "\n",
  "dP=", round(edges_df$dp_edge_Pa, 1), " 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 tree with index run highlighted (PDF)
## ---------------------------------------------------------------

p_tree <- 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",
    linewidth = 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",
    linewidth = 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 Random Flows, dP, and Index Run Highlighted")

ggsave("tree_flows_depth6_index_run.pdf", p_tree, width = 8, height = 11)
cat("Saved PDF: tree_flows_depth6_index_run.pdf in ", getwd(), "\n")

## ---------------------------------------------------------------
## 9. Build per-path summary (table data)
## ---------------------------------------------------------------

path_df <- do.call(
  rbind,
  lapply(seq_along(path_info), function(i) {
    x <- path_info[[i]]
    leaf_node <- tail(x$nodes, 1)
    data.frame(
      path_id  = i,
      leaf     = leaf_node,
      path_str = paste(x$nodes, collapse = " -> "),
      total_dP = x$total_dP,
      stringsAsFactors = FALSE
    )
  })
)

# Sort by descending total dP
path_df <- path_df[order(-path_df$total_dP), ]

cat("\nPer-path total dP summary:\n")
print(path_df)

## ---------------------------------------------------------------
## 10. Create PDF with table + chart of total dP per path
## ---------------------------------------------------------------

pdf("tree_paths_dp_summary.pdf", width = 8, height = 10)

# Page 1: table of paths and total dP
grid.newpage()
tbl <- tableGrob(
  path_df[, c("path_id", "leaf", "total_dP", "path_str")],
  rows = NULL
)
grid.draw(tbl)

# Page 2: bar chart of total dP per path (by leaf)
grid.newpage()
p_chart <- ggplot(path_df, aes(x = reorder(leaf, -total_dP), y = total_dP)) +
  geom_col(fill = "steelblue") +
  labs(
    title = "Total Pressure Drop per Path (root -> leaf)",
    x = "Leaf node",
    y = "Total dP (Pa)"
  ) +
  theme_minimal() +
  theme(axis.text.x = element_text(angle = 45, hjust = 1))

print(p_chart)

dev.off()

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