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

library(dodgr)
library(igraph)
library(ggplot2)

## ---------------------------------------------------------------
## 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 == "X"),
  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. Build edge dataframe with coordinates and labels
## ---------------------------------------------------------------

# Start from agg so we keep flow and dP directly
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 (merge can sometimes coerce to character)
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"
)

## ---------------------------------------------------------------
## 7. Plot with ggplot2 and save to PDF
## ---------------------------------------------------------------

p <- ggplot() +
  # edges
  geom_segment(
    data = edges_df,
    aes(x = x_from, y = y_from, xend = x_to, yend = y_to),
    arrow = arrow(length = unit(2, "mm")),
    lineend = "round",
    colour = "grey30"
  ) +
  # edge labels
  geom_text(
    data = edges_df,
    aes(x = x_mid, y = y_mid, label = label),
    size = 2.2,
    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 and Pressure Drops")

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

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