# =============================================================================
# manuscript-viz.R
# Publication helpers: justify the number of clusters (silhouette / elbow),
# visualise clusters in 2-D (PCA / UMAP / t-SNE), and compare clustering
# solutions to each other when no ground truth exists.
# =============================================================================

# distinct, colour-blind-friendly cluster palette (ColorPalette is a sequential
# ramp; for categorical cluster colours we want maximally distinct hues).
.mc_pal <- function(n) {
  base <- c("#4E79A7", "#F28E2B", "#E15759", "#76B7B2", "#59A14F",
            "#EDC948", "#B07AA1", "#9C755F", "#BAB0AC", "#D37295")
  if (n <= length(base)) base[seq_len(n)] else grDevices::hcl.colors(n, "Dark 3")
}

# coerce a fit / dist / list to an object-by-object distance matrix
.as_distm <- function(x, distmeasure = "euclidean") {
  if (is.list(x) && !is.null(x$DistM)) return(as.matrix(x$DistM))
  if (is.list(x) && is.null(x$DistM)) {                       # a List of sources
    Z <- do.call(cbind, lapply(x, function(m) { z <- scale(as.matrix(m)); z[is.na(z)] <- 0; z }))
    return(as.matrix(stats::dist(Z)))
  }
  as.matrix(x)
}

#' Justify the number of clusters (silhouette + elbow / scree)
#'
#' For a fused dissimilarity (or a fitted result, or a \code{List} of sources)
#' compute, over a range of \eqn{k}, the average silhouette width and the total
#' within-cluster dispersion (an elbow / scree curve), and plot them so the
#' chosen \eqn{k} can be justified in the manuscript.
#'
#' @param x A fitted result (with \code{$DistM}), a dissimilarity matrix, or a
#'   \code{List} of data sources (objects in rows).
#' @param k_range Integer vector of candidate cluster counts.
#' @param linkage Agglomeration method used to cut the dissimilarity.
#' @param select How to pick \code{best_k}: \code{"elbow"} (default) uses the
#'   knee of the within-cluster dispersion (scree) curve -- the \eqn{k} at which
#'   the curve starts to flatten; \code{"silhouette"} uses the maximum average
#'   silhouette width.
#' @param plot Logical; draw the two-panel figure.
#' @return Invisibly a data frame with \code{k}, \code{silhouette} and
#'   \code{within_dispersion}, plus attributes \code{"best_k"}, \code{"elbow_k"}
#'   and \code{"silhouette_k"}.
#' @seealso \code{\link{embedding_plot}}, \code{\link{SelectnrClusters}}
#' @examples
#' data(mosaic_toy)
#' fit <- M_ABCpp(mosaic_toy$List, numsim = 40, NC = 3)
#' cluster_selection_plot(fit, k_range = 2:6)
#' @export
cluster_selection_plot <- function(x, k_range = 2:8, linkage = "ward.D2",
                                   select = c("elbow", "silhouette"), plot = TRUE) {
  select <- match.arg(select)
  D  <- .as_distm(x); d <- stats::as.dist(D)
  hc <- fastcluster::hclust(d, method = linkage)
  res <- do.call(rbind, lapply(k_range, function(k) {
    cl  <- stats::cutree(hc, k)
    sil <- mean(cluster::silhouette(cl, d)[, "sil_width"])
    wss <- sum(vapply(split(seq_along(cl), cl), function(idx)
      if (length(idx) < 2) 0 else sum(as.matrix(D)[idx, idx][lower.tri(diag(length(idx)))]),
      numeric(1)))
    data.frame(k = k, silhouette = sil, within_dispersion = wss)
  }))
  # ---- elbow (kneedle): the k where the scree curve starts to flatten. Scale
  # k and dispersion to [0,1] and take the point furthest below the straight
  # line joining the first and last point (largest drop in marginal gain).
  kk <- res$k; w <- res$within_dispersion
  kn <- (kk - min(kk)) / (max(kk) - min(kk))
  wn <- (w  - min(w))  / (max(w)  - min(w) + 1e-12)
  # distance below the chord from (0,1) to (1,0): d = (kn + wn - 1)/sqrt(2), but
  # the elbow of a decreasing convex curve is where (chord - curve) is largest.
  chord <- 1 - kn                                   # straight line from wn(0)=1 to wn(1)=0
  elbow_k <- res$k[which.max(chord - wn)]
  sil_k   <- res$k[which.max(res$silhouette)]
  best    <- if (select == "elbow") elbow_k else sil_k
  if (isTRUE(plot)) {
    if (all(graphics::par("mfrow") == c(1, 1))) {
      op <- graphics::par(mfrow = c(1, 2), mar = c(4, 4, 3, 1)); on.exit(graphics::par(op))
    }
    plot(res$k, res$within_dispersion, type = "b", pch = 19, xlab = "number of clusters (k)",
         ylab = "within-cluster dispersion", main = "Elbow / scree (chosen k = dashed)")
    graphics::abline(v = elbow_k, lty = 2, col = "red")
    plot(res$k, res$silhouette, type = "b", pch = 19, xlab = "number of clusters (k)",
         ylab = "average silhouette width", main = "Silhouette (for reference)")
    graphics::abline(v = sil_k, lty = 3, col = "grey50")
  }
  attr(res, "best_k")     <- best
  attr(res, "elbow_k")    <- elbow_k
  attr(res, "silhouette_k") <- sil_k
  invisible(res)
}

#' Show the clustering at every k (see how it deteriorates)
#'
#' Rather than only reporting the best \eqn{k}, project the objects to a
#' \emph{fixed} 2-D embedding and recolour them by the \eqn{k}-cut for each
#' candidate \eqn{k} in a grid, annotating the average silhouette per panel. This
#' makes it visible how the structure splits and degrades as \eqn{k} grows.
#'
#' @param x A fitted result (\code{$DistM}), a dissimilarity, or a \code{List}.
#' @param k_range Candidate cluster counts to display.
#' @param linkage Agglomeration method used to cut the fused dissimilarity.
#' @param embed Embedding for the (fixed) 2-D coordinates: \code{"pca"} (default),
#'   \code{"umap"} or \code{"tsne"}.
#' @param ncol Columns in the panel grid.
#' @return Invisibly a named numeric vector of the average silhouette per \eqn{k}.
#' @seealso \code{\link{cluster_selection_plot}}, \code{\link{embedding_plot}}
#' @examples
#' data(mosaic_toy)
#' fit <- M_ABCpp(mosaic_toy$List, numsim = 40, NC = 3)
#' cluster_k_sweep(fit, k_range = 2:6)
#' @export
cluster_k_sweep <- function(x, k_range = 2:6, linkage = "ward.D2",
                            embed = c("pca", "umap", "tsne"), ncol = 3) {
  embed <- match.arg(embed)
  D <- .as_distm(x); d <- stats::as.dist(D); n <- nrow(D)
  hc <- fastcluster::hclust(d, method = linkage)
  Y <- switch(embed,
    pca  = stats::cmdscale(d, k = 2),
    umap = { if (!requireNamespace("uwot", quietly = TRUE)) stop("embed='umap' needs 'uwot'.")
             uwot::umap(d, n_neighbors = min(15, n - 1), n_components = 2) },
    tsne = { if (!requireNamespace("Rtsne", quietly = TRUE)) stop("embed='tsne' needs 'Rtsne'.")
             Rtsne::Rtsne(d, is_distance = TRUE, dims = 2,
                          perplexity = max(5, min(30, floor((n - 1) / 3))))$Y })
  op <- graphics::par(mfrow = c(ceiling(length(k_range) / ncol), ncol), mar = c(3, 3, 3, 1))
  on.exit(graphics::par(op))
  sils <- vapply(k_range, function(k) {
    cl  <- stats::cutree(hc, k)
    sil <- mean(cluster::silhouette(cl, d)[, "sil_width"])
    plot(Y, col = .mc_pal(k)[cl], pch = 19, cex = 0.8,
         xlab = paste0(toupper(embed), " 1"), ylab = paste0(toupper(embed), " 2"),
         main = sprintf("k = %d   (silhouette %.2f)", k, sil))
    sil
  }, numeric(1))
  invisible(stats::setNames(sils, k_range))
}

#' Two-dimensional cluster visualisation (PCA / UMAP / t-SNE)
#'
#' Project the objects to two dimensions and colour them by cluster. \code{"pca"}
#' uses classical MDS of the fused dissimilarity (always available); \code{"umap"}
#' needs \pkg{uwot} and \code{"tsne"} needs \pkg{Rtsne}.
#'
#' @param x A fitted result (\code{$DistM}), a dissimilarity, or a \code{List}.
#' @param labels Cluster label of each object.
#' @param method \code{"pca"} (default), \code{"umap"} or \code{"tsne"}.
#' @param col Optional cluster colours.
#' @param main Plot title.
#' @return Invisibly the 2-column coordinate matrix.
#' @seealso \code{\link{cluster_selection_plot}}
#' @examples
#' data(mosaic_toy)
#' fit <- M_ABCpp(mosaic_toy$List, numsim = 40, NC = 3)
#' embedding_plot(fit, mosaic_labels(fit, 3), method = "pca")
#' @export
embedding_plot <- function(x, labels, method = c("pca", "umap", "tsne"),
                           col = NULL, main = NULL) {
  method <- match.arg(method)
  D <- .as_distm(x); n <- nrow(D)
  Y <- switch(method,
    pca  = stats::cmdscale(stats::as.dist(D), k = 2),
    umap = { if (!requireNamespace("uwot", quietly = TRUE))
               stop("method = 'umap' needs the 'uwot' package.")
             uwot::umap(stats::as.dist(D), n_neighbors = min(15, n - 1), n_components = 2) },
    tsne = { if (!requireNamespace("Rtsne", quietly = TRUE))
               stop("method = 'tsne' needs the 'Rtsne' package.")
             Rtsne::Rtsne(as.dist(D), is_distance = TRUE, dims = 2,
                          perplexity = max(5, min(30, floor((n - 1) / 3))))$Y })
  cl <- factor(labels)
  if (is.null(col)) col <- .mc_pal(nlevels(cl))
  if (is.null(main)) main <- paste0(toupper(method), " coloured by cluster")
  plot(Y, col = col[as.integer(cl)], pch = 19, xlab = paste(toupper(method), "1"),
       ylab = paste(toupper(method), "2"), main = main)
  graphics::legend("topright", levels(cl), col = col[seq_len(nlevels(cl))],
                   pch = 19, bty = "n", title = "cluster", cex = 0.8)
  invisible(Y)
}

#' Visually compare clustering solutions with ComparePlot (no ground truth)
#'
#' Aligns several clustering solutions of the same objects and draws them with
#' \code{\link{ComparePlot}} so their agreement can be read off directly -- which
#' objects the methods place together and where they disagree -- without any
#' ground truth. Each solution may be a fitted result (with \code{$Clust} /
#' \code{$DistM}) or a plain label vector; label vectors are turned into a
#' co-membership dissimilarity so every solution is comparable.
#'
#' @param solutions A named list of fitted results and/or label vectors, all over
#'   the same objects in the same order.
#' @param nrclusters Number of clusters to cut each solution into.
#' @param cols Optional cluster colours; defaults to a distinct palette.
#' @param cex.labels,margins Passed to \code{\link{ComparePlot}}.
#' @param ... Further arguments to \code{\link{ComparePlot}}.
#' @return Invisibly \code{NULL}; called for the plot. Needs \pkg{circlize} and
#'   \pkg{plotrix}.
#' @seealso \code{\link{clustering_concordance}}, \code{\link{ComparePlot}}
#' @examples
#' data(mosaic_toy)
#' sols <- list(
#'   var    = M_ABCpp(mosaic_toy$List, numsim = 40, NC = 3, weighting = c("var", "var")),
#'   nugget = M_ABCpp(mosaic_toy$List, numsim = 40, NC = 3, weighting = c("nugget", "nugget")),
#'   NEMO   = NEMO(mosaic_toy$List, k = 3)$cluster)
#' \dontrun{compare_methods_plot(sols, nrclusters = 3)}
#' @export
compare_methods_plot <- function(solutions, nrclusters, cols = NULL,
                                 cex.labels = 0.5, margins = c(5, 11, 3, 2), ...) {
  if (!requireNamespace("circlize", quietly = TRUE) ||
      !requireNamespace("plotrix", quietly = TRUE))
    stop("compare_methods_plot() needs the 'circlize' and 'plotrix' packages.")
  # object names: from the first fit's DistM, else first labelled vector, else index
  first <- solutions[[1]]
  samp  <- if (is.list(first) && !is.null(first$DistM)) rownames(as.matrix(first$DistM))
           else names(first)
  n <- if (is.list(first) && !is.null(first$DistM)) nrow(as.matrix(first$DistM)) else length(first)
  if (is.null(samp)) samp <- sprintf("obj%03d", seq_len(n))

  # build a hierarchy + co-membership distance from *any* solution shape:
  # a fitted result with a usable $Clust, or a plain label vector, and fall back
  # to the object's labels (via mosaic_labels) if $Clust cannot be coerced.
  from_labels <- function(cl) {
    cl <- as.integer(factor(cl)); if (length(cl) != length(samp)) cl <- rep(1L, length(samp))
    D  <- outer(cl, cl, function(a, b) as.numeric(a != b)); dimnames(D) <- list(samp, samp)
    hc <- stats::as.hclust(cluster::agnes(stats::as.dist(D), method = "average"))
    hc$labels <- samp; list(hc = hc, D = D)
  }
  wrap <- function(x) {
    hc <- NULL; D <- NULL
    if (is.list(x) && !is.null(x$Clust)) {
      hc <- tryCatch(stats::as.hclust(x$Clust), error = function(e) NULL)
      if (!is.null(hc)) {
        D <- tryCatch(as.matrix(x$DistM), error = function(e) NULL)
        if (is.null(D) || length(D) == 0L || nrow(D) != length(samp)) {
          cl <- stats::cutree(hc, min(nrclusters, length(hc$order)))
          D  <- outer(cl, cl, function(a, b) as.numeric(a != b))
        }
        if (is.null(rownames(D))) dimnames(D) <- list(samp, samp)
        if (is.null(hc$labels)) hc$labels <- samp
      }
    }
    if (is.null(hc)) {                       # label vector, or a $Clust we could not coerce
      cl <- if (is.list(x)) tryCatch(mosaic_labels(x, nrclusters), error = function(e) NULL) else x
      fl <- from_labels(if (is.null(cl)) rep(1L, length(samp)) else cl)
      hc <- fl$hc; D <- fl$D
    }
    hc$order.lab <- hc$labels[hc$order]
    el <- list(Clust = hc, DistM = D); attr(el, "method") <- "Single Clustering"; el
  }
  res <- lapply(solutions, wrap)
  if (is.null(cols)) cols <- .mc_pal(nrclusters)
  ComparePlot(res, nrclusters = nrclusters, cols = cols, names = names(solutions),
              circle = FALSE, plottype = "sweave", cex.labels = cex.labels,
              margins = margins, ...)
  invisible(NULL)
}

#' Compare clustering solutions to each other (no ground truth needed)
#'
#' Given several labellings of the same objects (e.g. different weightings or
#' methods), compute their pairwise agreement (ARI) and each object's stability
#' -- how consistently it is co-clustered across the solutions. Useful for real
#' data where no ground truth is available.
#'
#' @param labels_list A named list of label vectors, all of the same length.
#' @return A list with \code{agreement} (pairwise ARI matrix), \code{stability}
#'   (mean pairwise co-membership per object) and \code{changed} (objects whose
#'   label is not constant across solutions).
#' @seealso \code{\link{cluster_agreement}}, \code{\link{compare_clusterings}}
#' @examples
#' data(mosaic_toy)
#' v <- mosaic_labels(M_ABCpp(mosaic_toy$List, numsim = 40, NC = 3,
#'                            weighting = c("var", "var")), 3)
#' n <- mosaic_labels(M_ABCpp(mosaic_toy$List, numsim = 40, NC = 3,
#'                            weighting = c("nugget", "nugget")), 3)
#' clustering_concordance(list(var = v, nugget = n))$agreement
#' @export
clustering_concordance <- function(labels_list) {
  stopifnot(length(labels_list) >= 2)
  nm <- names(labels_list); if (is.null(nm)) nm <- paste0("sol", seq_along(labels_list))
  m  <- length(labels_list); n <- length(labels_list[[1]])
  A  <- matrix(1, m, m, dimnames = list(nm, nm))
  for (i in seq_len(m)) for (j in seq_len(m)) if (i < j) {
    a <- cluster_agreement(labels_list[[i]], labels_list[[j]])[["ARI"]]
    A[i, j] <- A[j, i] <- a
  }
  # per-object stability: average pairwise co-membership agreement across solutions
  co <- matrix(0, n, n)
  for (l in labels_list) co <- co + outer(l, l, function(a, b) as.numeric(a == b))
  co <- co / m
  stability <- (rowSums(co) - 1) / (n - 1)
  changed <- which(apply(vapply(labels_list, function(l)
    match(l, unique(l)), integer(n)), 1, function(r) length(unique(r)) > 1))
  list(agreement = A, stability = stability, changed = changed)
}
