# source("1 PA Decline/data_format.R")

# NEW QUARTILES


# Visuals - sankey
# https://stackoverflow.com/questions/50395027/beautifying-sankey-alluvial-visualization-using-r


## Painting

df_raw <- readr::read_csv(here::here("2 Longterm/DDV 241031/event_sankey_data.csv")) 

df <- df_raw|> 
  dplyr::rename(pase_0_cut=pase_0_quartile,
                pase_6_cut=pase_4_quartile) |> 
  dplyr::mutate(change=dplyr::case_when(
    pase_0_cut==1 & pase_6_cut==1 ~ "ll",
    pase_0_cut %in% 2:4 & pase_6_cut %in% 2:4 ~ "hh",
    pase_0_cut<pase_6_cut ~ "hop",
    pase_0_cut>pase_6_cut ~ "drop"
  ),
  dplyr::across(c(pase_0_cut,pase_6_cut,change),as.factor))


df |> sankey_ready()

sankey_ready <- function(data,change.var="change"){
  df <- data 
  # |> 
  #   dplyr::count(dplyr::across(dplyr::all_of(c("pase_0_cut", "pase_6_cut",change.var)))) |> 
  #   dplyr::mutate(dplyr::across(dplyr::starts_with("pase_"),\(.x) factor(.x))) |> 
  #   setNames(c("pase_0_cut", "pase_6_cut","change","n"))
  
  lbs0 <-
    c(
      paste0("1st \n(n=", sum(df$n[df$pase_0_cut == "1"]), ")"),
      paste0("2nd \n(n=", sum(df$n[df$pase_0_cut == "2"]), ")"),
      paste0("3rd \n(n=", sum(df$n[df$pase_0_cut == "3"]), ")"),
      paste0("4th \n(n=", sum(df$n[df$pase_0_cut == "4"]), ")")
    )
  
  
  lbs6 <-
    c(
      paste0("1st \n(n=", sum(df$n[df$pase_6_cut == "1"]), ")"),
      paste0("2nd \n(n=", sum(df$n[df$pase_6_cut == "2"]), ")"),
      paste0("3rd \n(n=", sum(df$n[df$pase_6_cut == "3"]), ")"),
      paste0("4th \n(n=", sum(df$n[df$pase_6_cut == "4"]), ")")
    )
  
  
  levels(df$pase_0_cut) <- lbs0[1:length(levels(df$pase_0_cut))]
  levels(df$pase_6_cut) <- lbs6[1:length(levels(df$pase_6_cut))]
  
  df$pase_0_cut <- factor(df$pase_0_cut, levels = rev(levels(df$pase_0_cut)))
  df$pase_6_cut <- factor(df$pase_6_cut, levels = rev(levels(df$pase_6_cut)))
  
  df$change <- factor(df$change, levels = c("hh","hop", "drop", "ll"))
  
  if (change.var=="change"){
    df |> dplyr::mutate(first_grp=ifelse(substr(pase_0_cut,1,1)==1,"low","higher"))
  } else if (change.var=="change_any"){
    df |> dplyr::mutate(first_grp=dplyr::case_when(
      substr(pase_0_cut,1,1)==1 ~ "low",
      substr(pase_0_cut,1,1) %in% 2:3 ~ "mid",
      substr(pase_0_cut,1,1)==4 ~ "high"))
  }
}

# hops <- "#66c1a3" # grey
# # drops <- "#990033"  # Midtrød
# drops <- "#CE0045" # Lighter Midtrød
# nos <- "grey80" # Light grey
# 
# # border <- "#00596B"
# # box <- "#008099"
# 
# border <- "#EA571D"
# box <- "#1E4B66"
# 
# higher <- "yellow"
# low <- "purple"

library(ggalluvial)

library(ggplot2)

# project.aid::color_plot(viridisLite::turbo(4))

plot_sankey <- function(data,
                        # palette=viridisLite::turbo(4),
                        hops = "#1AE4B6FF",
                        drops = "#FABA39FF",
                        hh = "#30123BFF",
                        ll = "#7A0403FF",
                        border = "#EA571D",
                        box = "#1E4B66",
                        higher = "#1E4B66",
                        mid = "#1E4B66",
                        low = "#1E4B66",
                        alpha = 0.8,
                        a1=pase_0_cut,
                        a2=pase_6_cut,
                        a1.grp=first_grp,
                        text.size = 4
){
  
  if (length(unique(data[[ncol(data)]]))>2) {
    fills <- c(higher,low,mid)
  } else {
    fills <- c(higher,low) 
  }
  
  cls <- c(hh, hops, drops, ll)
  # stratum.grp <- c(df[["first_grp"]],df[["last_grp"]])
  
  # cls <- palette
  # browser()
  ggplot(data, aes(y = n, axis1 = {{a1}}, axis2 = {{a2}})) +
    geom_alluvium(
      aes(fill = change, color = change),
      width = 1 / 16,
      alpha = alpha,
      knot.pos = 0.4,
      curve_type ="sigmoid"
    ) +
    geom_stratum(aes(fill={{a1.grp}}),
                 # geom_stratum(aes(fill=stratum_grp),
                 size = 2, 
                 width = 1 / 3.4,
                 # fill = box,
                 color = border
    ) +
    geom_text(stat = "stratum", 
              aes(label = after_stat(stratum)), 
              colour = "white", 
              size = text.size,
              lineheight = 1) +
    scale_x_continuous(
      breaks = 1:2,
      labels = c("Pre-stroke\nPASE quartile", "Six months\nPASE quartile")
    ) +
    scale_fill_manual(values = c(cls,fills),na.value = box) +
    scale_color_manual(values = cls) +
    ggtitle("PA level changes from \npre-stroke to post-stroke")
}

# c("change","change_any") |> purrr::map(\(.x){
#   df_raw |>
#     sankey_ready(change.var = .x)
# }) |> 
#   purrr::map(\(.x){
#     .x |> plot_sankey(text.size=4.5)
#   }) |> 
#   patchwork::wrap_plots()

pal <- viridisLite::turbo(12)

project.aid::color_plot(pal)

df |>
  sankey_ready() |> 
  plot_sankey(text.size=4.5)

p_delta <- df |>
  sankey_ready() |> 
  plot_sankey(text.size=4.5)

ggplot2::ggsave(filename = here::here("2 Longterm/sankey_event_stroke.png"),
                p_delta +
                  theme_void() +
                  theme(
                    legend.position = "none",
                    # panel.grid.major = element_blank(),
                    # panel.grid.minor = element_blank(),
                    # axis.text.y = element_blank(),
                    # axis.title.y = element_blank(),
                    axis.text.x = element_text(),
                    # text = element_text(size = 5),
                    plot.title = element_blank(),
                    # panel.background = element_rect(fill = "white"),
                    plot.background = element_rect(fill="white"),
                    panel.border = element_blank()
                  ),
                units = "mm",
                width = 84,
                height = 70,
                # pointsize = 30,
                dpi = 600)