## Figures explaining what a column of a ComplexHeatmap density heatmap is.
## Reuses the TCGA-ESCA counts cached by normalize_tcga.R.
##
## Usage: Rscript density_explainer.R <work_dir> <figure_dir>

args   <- commandArgs(trailingOnly = TRUE)
WORK   <- if (length(args) >= 1) args[1] else "."
FIGDIR <- if (length(args) >= 2) args[2] else "."
CACHE  <- file.path(WORK, "esca_se.rds")

suppressPackageStartupMessages({
  library(SummarizedExperiment); library(edgeR); library(ComplexHeatmap)
  library(circlize); library(RColorBrewer); library(gifski)
})

se   <- readRDS(CACHE)
cts  <- assay(se, "unstranded")
info <- as.data.frame(colData(se))
grp  <- ifelse(info$sample_type == "Solid Tissue Normal", "Normal", "Tumor")
cts  <- cts[rowData(se)$gene_type == "protein_coding", ]
dge  <- DGEList(counts = cts, group = factor(grp))
dge  <- dge[filterByExpr(dge, group = factor(grp)), , keep.lib.sizes = FALSE]
M    <- log2(dge$counts + 1)                       # raw, the panel that shows the problem

YLIM <- c(-2, 16)
GRID <- seq(YLIM[1], YLIM[2], length.out = 512)

## density of every sample on one common grid, exactly what densityHeatmap does
dens_of <- function(x) {
  d <- density(x, na.rm = TRUE)
  approx(d$x, d$y, xout = GRID, rule = 2)$y
}
D <- apply(M, 2, dens_of)                          # 512 grid points x 44 samples

## the same color mapping densityHeatmap uses by default
col_fun <- colorRamp2(seq(0, quantile(D, 0.99, na.rm = TRUE), length = 11),
                      rev(brewer.pal(11, "Spectral")), space = "LAB")

TARGET <- 1L                                       # the sample we follow through
lab    <- sprintf("sample %d of %d", TARGET, ncol(M))

# ---------------------------------------------------------------------------
# 1. a conventional density curve, filled with the heatmap's own color scale
# ---------------------------------------------------------------------------
png(file.path(FIGDIR, "explain-01-curve.png"), width = 1400, height = 620, res = 150)
par(mar = c(4.2, 4.4, 3, 1))
d <- D[, TARGET]
plot(GRID, d, type = "n", xlab = "log2 expression", ylab = "density",
     main = sprintf("A density curve for one sample (%s)", lab), xaxs = "i", yaxs = "i",
     ylim = c(0, max(d) * 1.08), cex.main = 1)
for (i in seq_len(length(GRID) - 1)) {
  polygon(c(GRID[i], GRID[i+1], GRID[i+1], GRID[i]), c(0, 0, d[i+1], d[i]),
          col = col_fun(mean(d[i:(i+1)])), border = NA)
}
lines(GRID, d, lwd = 2, col = "gray20")
mtext("color = height of the curve at that point", side = 3, line = 0.1, cex = 0.8, col = "gray35")
dev.off(); cat("wrote explain-01-curve.png\n")

# ---------------------------------------------------------------------------
# 2. rotate it upright, then flatten it into a single column
# ---------------------------------------------------------------------------
png(file.path(FIGDIR, "explain-02-flatten.png"), width = 1500, height = 700, res = 150)
layout(matrix(1:3, nrow = 1), widths = c(1.25, 1.25, 1))
dmax <- max(d)

par(mar = c(4.2, 4.4, 3.4, 1))
plot(d, GRID, type = "n", xlab = "density", ylab = "log2 expression",
     main = "1. turn it on its side", ylim = YLIM, xlim = c(0, dmax * 1.05),
     xaxs = "i", yaxs = "i", cex.main = 1)
for (i in seq_len(length(GRID) - 1))
  polygon(c(0, d[i], d[i+1], 0), c(GRID[i], GRID[i], GRID[i+1], GRID[i+1]),
          col = col_fun(mean(d[i:(i+1)])), border = NA)
lines(d, GRID, lwd = 2, col = "gray20")

par(mar = c(4.2, 4.4, 3.4, 1))
plot(NA, xlim = c(0, 1), ylim = YLIM, xlab = "", ylab = "", xaxt = "n",
     main = "2. stretch every slice\nto the same width", yaxs = "i", cex.main = 1)
for (i in seq_len(length(GRID) - 1))
  rect(0, GRID[i], 1, GRID[i+1], col = col_fun(mean(d[i:(i+1)])), border = NA)

par(mar = c(4.2, 2, 3.4, 1))
plot(NA, xlim = c(0, 1), ylim = YLIM, xlab = "", ylab = "", xaxt = "n", yaxt = "n",
     main = "3. that is one column", yaxs = "i", cex.main = 1)
for (i in seq_len(length(GRID) - 1))
  rect(0.3, GRID[i], 0.7, GRID[i+1], col = col_fun(mean(d[i:(i+1)])), border = NA)
rect(0.3, YLIM[1], 0.7, YLIM[2], border = "red", lwd = 2)
dev.off(); cat("wrote explain-02-flatten.png\n")

# ---------------------------------------------------------------------------
# 3. the same column, outlined inside the real heatmap
# ---------------------------------------------------------------------------
png(file.path(FIGDIR, "explain-03-incontext.png"), width = 1500, height = 780, res = 150)
ht <- densityHeatmap(M, ylim = YLIM, ylab = "log2 expression",
                     title = "Every column is one sample's density curve, seen from above",
                     title_gp = gpar(fontsize = 12, fontface = "bold"),
                     cluster_columns = FALSE, show_column_names = FALSE)
draw(ht)
## densityHeatmap suffixes its name (e.g. "density_mKSwPdl5"), so ask the object
decorate_heatmap_body(ht@name, {
  n <- ncol(M)
  grid.rect(x = (TARGET - 0.5) / n, width = 1 / n, just = "center",
            gp = gpar(col = "red", lwd = 2.5, fill = NA))
})
dev.off(); cat("wrote explain-03-incontext.png\n")
# ---------------------------------------------------------------------------
# 4. animation: each sample is a real solid. Its profile is the density curve
#    and it is extruded sideways to give it thickness. The camera yaws and
#    pitches from an oblique view to straight overhead, where the top face of
#    each solid is the heatmap column: looking down the density axis, every
#    slice projects to the same footprint however tall it is, so only its
#    color survives.
# ---------------------------------------------------------------------------
NS   <- 10L
sel  <- round(seq(1, ncol(M), length.out = NS))
GA   <- seq(YLIM[1], YLIM[2], length.out = 110)
Ds   <- apply(M[, sel, drop = FALSE], 2, function(x) {
  dd <- density(x, na.rm = TRUE); approx(dd$x, dd$y, xout = GA, rule = 2)$y })

## put everything in comparable world units: a 10 x 10 table carrying 3-unit ridges
YW   <- (GA - YLIM[1]) / diff(YLIM) * 10
ZS   <- 3.0 / max(Ds)
HALF <- 0.40                                        # half-thickness of each slab

## yaw about the density axis, then pitch the table toward the viewer.
## alpha = 0 and theta = 0 is the straight-overhead view, i.e. the heatmap.
proj <- function(x, y, z, alpha, theta) {
  ct <- cos(theta); st <- sin(theta)
  x1 <- x * ct - y * st
  y1 <- x * st + y * ct
  list(sx = x1,
       sy = y1 * cos(alpha) + z * sin(alpha),
       sz = -y1 * sin(alpha) + z * cos(alpha))     # larger = nearer the viewer
}

A_MAX <- 62 * pi / 180
T_MAX <- 22 * pi / 180

## fixed limits over the whole rotation so the object does not drift
gx <- c(0.5 - HALF, NS + 0.5 + HALF); gy <- c(0, 10); gz <- c(0, 3.0)
gp <- expand.grid(x = gx, y = gy, z = gz, k = seq(0, 1, length.out = 20))
pp <- proj(gp$x, gp$y, gp$z, A_MAX * gp$k, T_MAX * gp$k)
XLIM <- range(pp$sx) + c(-0.3, 0.3); SYLIM <- range(pp$sy) + c(-0.5, 0.5)

shade <- function(hex, f) { v <- col2rgb(hex) / 255; rgb(v[1]*f, v[2]*f, v[3]*f) }

draw_frame <- function(k) {
  alpha <- A_MAX * k; theta <- T_MAX * k
  par(mar = c(2.0, 3.0, 3.2, 1))
  plot(NA, xlim = XLIM, ylim = SYLIM, xaxs = "i", yaxs = "i",
       xlab = "", ylab = "", xaxt = "n", yaxt = "n", bty = "n", asp = 1,
       main = "Each sample is a solid; the top face is the heatmap column",
       cex.main = 1.02)

  quads <- vector("list", NS * (length(GA) - 1) * 3); n <- 0L
  for (s in seq_len(NS)) {
    dz <- Ds[, s] * ZS
    for (i in seq_len(length(GA) - 1)) {
      y1 <- YW[i]; y2 <- YW[i + 1]; z1 <- dz[i]; z2 <- dz[i + 1]
      cc <- col_fun(mean(Ds[i:(i + 1), s]))
      ## top face: the surface that becomes the column
      p <- proj(c(s - HALF, s + HALF, s + HALF, s - HALF), c(y1, y1, y2, y2),
                c(z1, z1, z2, z2), alpha, theta)
      n <- n + 1L; quads[[n]] <- list(x = p$sx, y = p$sy, col = cc, d = mean(p$sz))
      ## the two side walls, which only have area once the scene is yawed
      if (k > 0.02) {
        for (sgn in c(1, -1)) {
          p <- proj(rep(s + sgn * HALF, 4), c(y1, y1, y2, y2), c(0, z1, z2, 0),
                    alpha, theta)
          n <- n + 1L
          quads[[n]] <- list(x = p$sx, y = p$sy,
                             col = shade(cc, if (sgn > 0) 0.62 else 0.45),
                             d = mean(p$sz))
        }
      }
    }
    ## near end cap, so the slab reads as solid rather than hollow
    if (k > 0.02) {
      p <- proj(c(s - HALF, s + HALF, s + HALF, s - HALF), rep(YW[1], 4),
                c(0, 0, Ds[1, s] * ZS, Ds[1, s] * ZS), alpha, theta)
      n <- n + 1L
      quads[[n]] <- list(x = p$sx, y = p$sy, col = shade(col_fun(Ds[1, s]), 0.5),
                         d = mean(p$sz))
    }
  }
  quads <- quads[seq_len(n)]
  ord <- order(vapply(quads, function(q) q$d, numeric(1)))   # painter's algorithm
  for (q in quads[ord]) polygon(q$x, q$y, col = q$col, border = NA)

  ## floor outline, so the rotation is readable
  fl <- proj(c(gx[1], gx[2], gx[2], gx[1]), c(gy[1], gy[1], gy[2], gy[2]),
             rep(0, 4), alpha, theta)
  polygon(fl$sx, fl$sy, border = rgb(0.45, 0.45, 0.45, 0.5), lwd = 1)

  mtext("samples", side = 1, line = 0.4, cex = 0.9)
  lab <- if (k > 0.9) "density is height" else
         if (k < 0.1) "from overhead, height is only color" else ""
  if (nzchar(lab)) mtext(lab, side = 3, line = 0.05, cex = 0.85, col = "gray30")
}

ease <- function(t) (1 - cos(pi * t)) / 2
FR   <- file.path(FIGDIR, "frames"); dir.create(FR, showWarnings = FALSE)
ks   <- c(rep(1, 4), 1 - ease(seq(0, 1, length.out = 18)), rep(0, 5))
ks   <- c(ks, rev(ks))
files <- character(length(ks))
for (j in seq_along(ks)) {
  files[j] <- sprintf("%s/fr-%03d.png", FR, j)
  png(files[j], width = 700, height = 460, res = 78); draw_frame(ks[j]); dev.off()
}
gifski(files, gif_file = file.path(FIGDIR, "explain-04-rotate.gif"),
       width = 700, height = 460, delay = 1/15, progress = FALSE)
unlink(FR, recursive = TRUE)
cat("wrote explain-04-rotate.gif (", length(ks), "frames )\n")
cat("done\n")