---
title: "LUCID Early Tutorial: Sankey Plot with Balanced Top Omics"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{LUCID Early Tutorial: Sankey Plot with Balanced Top Omics}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

## 1) Goal

A LUCID fit's Sankey diagram (`plot()`, in the other vignettes) shows the
package's own default view of exposure -> cluster -> outcome. This tutorial
builds a **custom** Sankey view on top of a fitted model for a situation
`plot()` doesn't handle out of the box: multiple named omics *families*
concatenated into one early-integration matrix (here, microbiota and
metabolite features), where a meaningful figure needs to show a *balanced*
number of top features from each family side by side, not just the overall
top-`N` (which could easily be dominated by one family). It walks through:

1. Simulating two omics families -- microbiota (`X0000001`, `X0000002`, ...)
   and metabolites (`hilic0001`, `c180001`, ...) -- concatenated into one
   `Z` matrix for early integration.
2. Fitting a LUCID early model with `estimate_lucid()`.
3. Ranking and selecting a balanced top-`N` per family, and plotting the
   result as a Sankey diagram.
4. Contrasting that dataset-specific ranking with the package's
   general-purpose `get_top_omics_features()`.

## 2) Setup

```{r setup, message=FALSE, warning=FALSE}
required_pkgs <- c("plotly")
missing_pkgs <- required_pkgs[!vapply(required_pkgs, requireNamespace, logical(1), quietly = TRUE)]
if (length(missing_pkgs) > 0) {
  stop("Please install required packages first: ", paste(missing_pkgs, collapse = ", "))
}

library(LUCIDus)
library(plotly)  # for %>% and plot_ly(), used directly by this vignette's own Sankey code

data(sim_data)
```

## 3) Simulate a Two-Layer Omics Matrix for Early Integration

As in the other vignettes, the simulation deliberately builds in a signal:
a latent state (`latent_signal`) shifts a subset of features in each family
apart by cluster, on top of a shared dependence on one exposure
(`g_anchor`), so the fitted model below has real structure to find and rank.
The two metabolite sub-blocks (`hilic`, `c18`) are concatenated into one
`metabolites` block before joining `microbiota` into the single `Z` matrix
early integration expects.

```{r simulate-omics}
set.seed(1008)

n <- 200

G <- sim_data$G[seq_len(n), , drop = FALSE]
colnames(G) <- paste0("exposure", seq_len(ncol(G)))

Y <- as.numeric(sim_data$Y_normal[seq_len(n), 1])
CoY <- sim_data$Covariate[seq_len(n), , drop = FALSE]
colnames(CoY) <- paste0("covariate", seq_len(ncol(CoY)))

latent_signal <- ifelse(sim_data$X[seq_len(n)] == 1, -1, 1)
g_anchor <- as.numeric(scale(G[, 1]))

simulate_omics_block <- function(feature_names, signal_n, latent_signal, g_anchor,
                                 cluster_scale, noise_sd, offset = 0) {
  n_obs <- length(latent_signal)
  p <- length(feature_names)
  block <- matrix(NA_real_, nrow = n_obs, ncol = p)

  signal_loadings <- seq(cluster_scale, 0.8, length.out = signal_n)

  for (j in seq_len(p)) {
    direction <- if (j %% 2L == 0L) 1 else -1
    cluster_loading <- if (j <= signal_n) {
      direction * signal_loadings[j]
    } else {
      runif(1, -0.20, 0.20)
    }
    exposure_loading <- if (j <= signal_n) {
      direction * runif(1, 0.15, 0.35)
    } else {
      runif(1, -0.10, 0.10)
    }

    block[, j] <- offset +
      cluster_loading * latent_signal +
      exposure_loading * g_anchor +
      rnorm(n_obs, sd = noise_sd)
  }

  block <- scale(block)
  colnames(block) <- feature_names
  block
}

n_microbiota <- 18
n_hilic <- 12
n_c18 <- 8

microbiota_names <- paste0("X", sprintf("%07d", seq_len(n_microbiota)))
hilic_names <- paste0("hilic", sprintf("%04d", seq_len(n_hilic)))
c18_names <- paste0("c18", sprintf("%04d", seq_len(n_c18)))

microbiota <- simulate_omics_block(
  feature_names = microbiota_names,
  signal_n = 6,
  latent_signal = latent_signal,
  g_anchor = g_anchor,
  cluster_scale = 1.80,
  noise_sd = 0.85
)

hilic <- simulate_omics_block(
  feature_names = hilic_names,
  signal_n = 5,
  latent_signal = latent_signal,
  g_anchor = g_anchor,
  cluster_scale = 1.55,
  noise_sd = 0.90,
  offset = 0.15
)

c18 <- simulate_omics_block(
  feature_names = c18_names,
  signal_n = 4,
  latent_signal = latent_signal,
  g_anchor = g_anchor,
  cluster_scale = 1.35,
  noise_sd = 0.95,
  offset = -0.10
)

metabolites <- cbind(hilic, c18)
Z <- cbind(microbiota, metabolites)

omics_dimensions <- data.frame(
  block = c("microbiota", "hilic", "c18", "metabolites_total", "combined_Z"),
  n_features = c(ncol(microbiota), ncol(hilic), ncol(c18), ncol(metabolites), ncol(Z))
)
omics_dimensions
```

`combined_Z`'s feature count is the sum of `microbiota` and
`metabolites_total`, confirming the concatenation produced the single matrix
early integration expects -- one joint cluster variable will be fit from all
38 features together, not one per family.

## 4) Fit a LUCID Early Model

`estimate_lucid(..., lucid_model = "early")` fits one latent cluster
variable, `K = 2` here, from the pooled `G`/`Z`/`Y`. `res_Mu` (`K x M`) holds
the fitted cluster-specific mean for every omics feature -- the ranking used
throughout this tutorial reads directly from this matrix.

```{r fit-early, message=FALSE}
fit_early <- estimate_lucid(
  G = G,
  Z = Z,
  Y = Y,
  lucid_model = "early",
  CoY = CoY,
  family = "normal",
  K = 2,
  seed = 1008,
  max_itr = 20,
  max_tot.itr = 60,
  tol = 1e-2
)

class(fit_early)
dim(fit_early$res_Mu)
```

`dim(fit_early$res_Mu)` should read `2 x 38`: one row per cluster, one
column per omics feature (the same 38 confirmed above).

## 5) Sankey Function with Equal Top N per Omics Family

The Sankey below has three link stages -- exposure -> cluster,
cluster -> omics feature, cluster -> outcome -- built directly from the
fitted model's own parameters (`res_Beta` for the first, `res_Mu` for the
second, `res_Gamma` for the third); no new estimation happens here, only
reshaping and coloring for the plot. `select_top_omics_balanced()` is the
piece that makes the "balanced" part work: it scores every feature by how
far apart its two cluster means are (`abs(mu[c1,] - mu[c2,])`) and keeps the
top `top_n_each` **within each omics family separately**, so a family with
many features can't crowd out a smaller one just by having more entries to
rank.

```{r sankey-function}
detect_omics_family <- function(x) {
  out <- rep("exposure", length(x))
  out[grepl("^Latent Cluster", x)] <- "lc"
  out[grepl("^X\\d{7}$", x)] <- "microbiota"
  out[grepl("^(hilic|c18)", x, ignore.case = TRUE)] <- "metabolite"
  out
}

select_top_omics_balanced <- function(mu_mat, z_names, cluster_pair, top_n_each = NULL) {
  if (is.null(top_n_each)) {
    return(z_names)
  }

  if (!is.numeric(top_n_each) || length(top_n_each) != 1 || is.na(top_n_each) || top_n_each < 1) {
    stop("`top_n_each` must be a single positive integer or NULL.")
  }

  top_n_each <- as.integer(top_n_each)
  c1 <- as.integer(cluster_pair[1])
  c2 <- as.integer(cluster_pair[2])
  score <- abs(mu_mat[c1, ] - mu_mat[c2, ])
  family <- detect_omics_family(z_names)

  ordered_keep <- character(0)
  for (fam in c("microbiota", "metabolite")) {
    fam_idx <- which(family == fam)
    if (!length(fam_idx)) {
      next
    }
    fam_idx <- fam_idx[order(score[fam_idx], decreasing = TRUE)]
    ordered_keep <- c(
      ordered_keep,
      z_names[fam_idx[seq_len(min(top_n_each, length(fam_idx)))]]
    )
  }

  ordered_keep
}

make_link_color <- function(target_group, is_positive) {
  pos_cols <- c(
    lc = "#8E9BAE",
    microbiota = "#1B7F5A",
    metabolite = "#C46A1B",
    outcome = "#B23A48",
    exposure = "#5E6F86"
  )
  neg_cols <- c(
    lc = "#DCE3EA",
    microbiota = "#BFE2D6",
    metabolite = "#F2D4B5",
    outcome = "#E9C0C6",
    exposure = "#CCD6E2"
  )

  ifelse(is_positive, pos_cols[target_group], neg_cols[target_group])
}

sankey_early_integration <- function(
    lucid_fit1,
    top_n_each = 5,
    cluster_pair = c(1, 2),
    text_size = 13,
    drop_covariates = TRUE
) {
  if (!inherits(lucid_fit1, "early_lucid")) {
    stop("`lucid_fit1` must inherit from class 'early_lucid'.")
  }

  needed <- c("K", "var.names", "res_Beta", "res_Mu", "res_Gamma")
  miss <- needed[!needed %in% names(lucid_fit1)]
  if (length(miss) > 0) {
    stop("Missing fields in `lucid_fit1`: ", paste(miss, collapse = ", "))
  }

  K <- lucid_fit1$K
  var.names <- lucid_fit1$var.names
  Gnames <- var.names$Gnames
  Znames <- var.names$Znames
  Yname <- var.names$Ynames

  beta_mat <- as.matrix(lucid_fit1$res_Beta)
  mu_mat <- as.matrix(lucid_fit1$res_Mu)

  if (nrow(beta_mat) != K) stop("Unexpected `res_Beta` shape.")
  if (nrow(mu_mat) != K) stop("Unexpected `res_Mu` shape.")

  if (length(cluster_pair) != 2) stop("`cluster_pair` must be length-2.")
  c1 <- as.integer(cluster_pair[1])
  c2 <- as.integer(cluster_pair[2])
  if (c1 < 1 || c1 > K || c2 < 1 || c2 > K || c1 == c2) {
    stop("`cluster_pair` must reference two distinct cluster indices in 1..K.")
  }

  keep_z <- select_top_omics_balanced(
    mu_mat = mu_mat,
    z_names = Znames,
    cluster_pair = cluster_pair,
    top_n_each = top_n_each
  )

  valueGtoX <- as.vector(t(beta_mat[, -1, drop = FALSE]))

  mu_use <- mu_mat[, Znames %in% keep_z, drop = FALSE]
  mu_use <- mu_use[, keep_z, drop = FALSE]
  z_use <- colnames(mu_use)
  valueXtoZ <- as.vector(t(mu_use))

  gamma_beta <- if (!is.null(lucid_fit1$res_Gamma$beta)) {
    as.numeric(lucid_fit1$res_Gamma$beta)
  } else {
    as.numeric(lucid_fit1$res_Gamma$fit$coefficients)
  }
  if (length(gamma_beta) < K) stop("Could not extract cluster->outcome effects.")

  valueXtoY <- gamma_beta[seq_len(K)]
  valueXtoY[1] <- 0

  GtoX <- data.frame(
    source = rep(Gnames, K),
    target = paste0("Latent Cluster", rep(seq_len(K), each = length(Gnames))),
    value = abs(valueGtoX),
    group = valueGtoX > 0,
    stringsAsFactors = FALSE
  )

  XtoZ <- data.frame(
    source = paste0("Latent Cluster", rep(seq_len(K), each = length(z_use))),
    target = rep(z_use, K),
    value = abs(valueXtoZ),
    group = valueXtoZ > 0,
    stringsAsFactors = FALSE
  )

  XtoY <- data.frame(
    source = paste0("Latent Cluster", seq_len(K)),
    target = rep(Yname, K),
    value = abs(valueXtoY),
    group = valueXtoY > 0,
    stringsAsFactors = FALSE
  )

  links <- rbind(GtoX, XtoZ, XtoY)

  if (drop_covariates) {
    links <- links[!grepl("^covariate", links$source, ignore.case = TRUE), , drop = FALSE]
  }

  nodes <- data.frame(name = unique(c(links$source, links$target)), stringsAsFactors = FALSE)
  nodes$group <- detect_omics_family(nodes$name)
  nodes$group[nodes$name %in% Yname] <- "outcome"

  group_order <- c("exposure", "lc", "microbiota", "metabolite", "outcome")
  nodes$group <- factor(nodes$group, levels = group_order)
  nodes <- nodes[order(nodes$group, nodes$name), , drop = FALSE]
  nodes$group <- as.character(nodes$group)

  node_color_map <- c(
    exposure = "#4C6A92",
    lc = "#B9C6D3",
    microbiota = "#2A9D6F",
    metabolite = "#D98C3F",
    outcome = "#C44536"
  )
  nodes$color <- unname(node_color_map[nodes$group])

  links$IDsource <- match(links$source, nodes$name) - 1L
  links$IDtarget <- match(links$target, nodes$name) - 1L

  target_group <- nodes$group[match(links$target, nodes$name)]
  target_group[grepl("^Latent Cluster", links$target)] <- "lc"
  links$link_color <- make_link_color(target_group, links$group)

  x_map <- c(
    exposure = 0.02,
    lc = 0.34,
    microbiota = 0.68,
    metabolite = 0.80,
    outcome = 0.97
  )
  nodes$x <- unname(x_map[nodes$group])

  nodes$y <- NA_real_
  for (g in group_order) {
    idx <- which(nodes$group == g)
    if (!length(idx)) {
      next
    }
    nodes$y[idx] <- if (length(idx) == 1) 0.5 else seq(0.03, 0.97, length.out = length(idx))
  }

  plotly::plot_ly(
    type = "sankey",
    orientation = "h",
    node = list(
      label = nodes$name,
      color = nodes$color,
      pad = 16,
      thickness = 18,
      line = list(color = "black", width = 0.35),
      x = nodes$x,
      y = nodes$y
    ),
    link = list(
      source = links$IDsource,
      target = links$IDtarget,
      value = pmax(links$value, 1e-12),
      color = links$link_color
    )
  ) %>%
    plotly::layout(
      font = list(size = text_size),
      margin = list(l = 20, r = 20, t = 30, b = 20)
    )
}
```

## 6) Top Features Selected per Family

```{r selected-features}
selected_top <- select_top_omics_balanced(
  mu_mat = fit_early$res_Mu,
  z_names = fit_early$var.names$Znames,
  cluster_pair = c(1, 2),
  top_n_each = 5
)

data.frame(
  feature = selected_top,
  family = detect_omics_family(selected_top),
  stringsAsFactors = FALSE
)
```

`select_top_omics_balanced()` above is specific to this dataset: it scores
features by their difference between one *pair* of clusters and balances the
count kept from each omics family by name pattern. The package's
general-purpose ranking function, `get_top_omics_features()`, instead scores
every feature by its separation across *all* clusters (the same criterion
`plot_cluster_omic_profile()` uses) with no family balancing or pairwise
comparison:

```{r top-omics-features}
top_general <- get_top_omics_features(fit_early, top_n = 10)
top_general
```

## 7) Display the Plot

Read either figure left to right: exposure nodes flow into the two latent
cluster nodes (link width proportional to `|G -> X coefficient|`, color by
sign), the clusters flow into omics feature nodes (width by
`|cluster mean|`), and the clusters flow into the outcome node. A feature
whose incoming link is wide is one the model estimates a large cluster-mean
difference for; a thin link means that feature barely differs between
clusters despite appearing in the diagram.

### 7.1 All omics in the concatenated early-integration matrix

`top_n_each = NULL` keeps every one of the 38 features, so this is a
complete but visually dense view -- useful as a first look, less useful for
picking out which features actually matter.

```{r plot-all, fig.width=13, fig.height=9}
fig_all <- sankey_early_integration(
  lucid_fit1 = fit_early,
  top_n_each = NULL,
  cluster_pair = c(1, 2),
  text_size = 11
)
fig_all
```

### 7.2 Top 5 microbiota and top 5 metabolites

Setting `top_n_each = 5` trims the diagram to the 5 most cluster-differentiating
features from *each* family (10 total omics nodes), which is where the
"balanced" ranking in section 5 matters: without it, a simple overall top-10
could plausibly be all microbiota (or all metabolites) if one family's
effects happen to be larger in scale, hiding the other family's real signal
entirely from the figure.

```{r plot-top, fig.width=13, fig.height=8}
fig_top_balanced <- sankey_early_integration(
  lucid_fit1 = fit_early,
  top_n_each = 5,
  cluster_pair = c(1, 2),
  text_size = 13
)
fig_top_balanced
```

## 8) Notes

- `Z` is a single early-integration matrix, but it is built by concatenating microbiota and metabolite blocks.
- Microbiota and metabolites can have different dimensions; `top_n_each` still keeps the same count from each family.
- Node colors are distinct for exposure, latent cluster, microbiota, metabolite, and outcome.
- Ranking is based on `abs(Mu[cluster_pair[1], ] - Mu[cluster_pair[2], ])` within each omics family.
- `select_top_omics_balanced()` and `sankey_early_integration()` are
  dataset-specific helpers defined in this vignette (section 5), not part of
  the package -- for a general-purpose, family-agnostic top-feature ranking
  on any fitted model, use the package's own `get_top_omics_features()`
  (section 6).