# Load required package
library(DiagrammeR)
library(dplyr)

# Vectorised function for asterisk notation
add_stars <- function(p) {
  case_when(
    is.na(p)        ~ "",
    p < 0.001       ~ "***",
    p < 0.01        ~ "**",
    p < 0.05        ~ "*",
    TRUE            ~ ""
  )
}

# Choose predictor 
selected_predictor <- "CFRS" 
predictor_node <- selected_predictor
predictor_label <- "Functional Cognitive Impairment"  

# Filter and annotate data
df_pred <- Path_Coefficients_Unadjusted %>%
  filter(Predictor == selected_predictor) %>%
  mutate(
    # Human-readable outcome node names
    Outcome_Node = paste0("Outcome_", Outcome),
    a_label = paste0(sprintf("%.2f", Path_A_Beta), add_stars(Path_A_p)),
    b_label = paste0(sprintf("%.2f", Path_B_Beta), add_stars(Path_B_p))
  )

# Node label mappings (Mediators)
mediator_labels <- df_pred %>%
  distinct(Mediator) %>%
  mutate(
    Mediator_Label = case_when(
      Mediator == "Active_Approach" ~ "Actively Approach",
      Mediator == "Withdrawal" ~ "Withdrawal",
      Mediator == "Social_Support" ~ "Seeking Social Support",
      Mediator == "Self_blame" ~ "Self-blame", 
      Mediator == "Planning" ~ "Refocus on Planning",
      Mediator == "Reappraisal" ~ "Positive Reappraisal",
      TRUE ~ Mediator
    )
  )

# Node label mappings (Outcomes)
outcome_labels <- df_pred %>%
  distinct(Outcome_Node, Outcome) %>%
  mutate(
    Outcome_Label = case_when(
      Outcome == "WEMWBS" ~ "Mental Wellbeing",
      Outcome == "GDS" ~ "Depression",
      Outcome == "GAD" ~ "Anxiety",
      TRUE ~ Outcome
    )
  )

# Edges (a and b paths)
a_edges <- df_pred %>%
  distinct(Predictor, Mediator, a_label) %>%
  mutate(edge = paste0(Predictor, " -> ", Mediator, ' [label = "', a_label, '"];')) %>%
  pull(edge)

b_edges <- df_pred %>%
  distinct(Mediator, Outcome_Node, b_label) %>%
  mutate(edge = paste0(Mediator, " -> ", Outcome_Node, ' [label = "', b_label, '"];')) %>%
  pull(edge)

all_edges <- paste(c(a_edges, b_edges), collapse = "\n  ")

# Node declarations
predictor_node_def <- paste0(predictor_node, ' [label = "', predictor_label, '", shape = box, style = filled, fillcolor = lightblue]')

mediator_nodes <- paste0(
  mediator_labels$Mediator, 
  ' [label = "', mediator_labels$Mediator_Label, '", shape = ellipse, style = filled, fillcolor = lightyellow]'
) %>% paste(collapse = "\n  ")

outcome_nodes <- paste0(
  outcome_labels$Outcome_Node, 
  ' [label = "', outcome_labels$Outcome_Label, '", shape = box, style = filled, fillcolor = lightpink]'
) %>% paste(collapse = "\n  ")

# Graphviz diagram code
diagram_code <- paste0("
digraph ", predictor_node, "_mediation {
  graph [layout = dot, rankdir = LR]

  ", predictor_node_def, "
  ", mediator_nodes, "
  ", outcome_nodes, "

  ", all_edges, "
}
")

# Render diagram
grViz(diagram_code)