## ----setup, include=FALSE-----------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment  = "#>",
  fig.width  = 7,
  fig.height = 4.5,
  dpi = 110
)
# Smooth vignette chunks need mgcv (Suggests) and CmdStan.
run_smooth_vignette <-
  instantiate::stan_cmdstan_exists()

## ----show-formulae, eval=FALSE------------------------------------------------
# ~ type + s(pseudotime, k = 5)

## ----load-libs, message=FALSE, warning=FALSE----------------------------------
library(dplyr)
library(tidyr)
library(ggplot2)
library(sccomp)
data("counts_obj")

## ----make-pseudotime----------------------------------------------------------
set.seed(1)
samples <- levels(counts_obj$sample)
pt_map <- tibble(
  sample     = samples,
  pseudotime = seq(0, 6, length.out = length(samples)) +
                 rnorm(length(samples), 0, 0.05)
)

counts_pt <- counts_obj |> left_join(pt_map, by = "sample")
counts_pt |> distinct(sample, type, pseudotime) |> head()

## ----fit-smooth, message=FALSE, warning=FALSE, eval = run_smooth_vignette-----
# fit <- counts_pt |>
#   sccomp_estimate(
#     formula_composition = ~ type + s(pseudotime, k = 5),
#     formula_variability = ~ 1,
#     sample              = "sample",
#     cell_group          = "cell_group",
#     abundance           = "count",
#     cores               = 1,
#     verbose             = FALSE,
#     max_sampling_iterations = 2000
#   )

## ----test-table, eval = run_smooth_vignette-----------------------------------
# fit |>
#   sccomp_test() |>
#   select(cell_group, parameter, c_lower, c_effect, c_upper, c_FDR)

## ----predict-grid, message=FALSE, warning=FALSE, eval = run_smooth_vignette----
# grid <- expand.grid(
#   pseudotime = seq(min(pt_map$pseudotime), max(pt_map$pseudotime), length.out = 60),
#   type       = levels(counts_obj$type)
# ) |>
#   as_tibble() |>
#   mutate(sample = sprintf("grid_%03d", row_number()))
# 
# pred <- fit |>
#   sccomp_predict(new_data = grid, number_of_draws = 200)
# 
# head(pred)

## ----plot-curves, fig.width=8, fig.height=5, eval = run_smooth_vignette-------
# top_cgs <- pred |>
#   group_by(cell_group) |>
#   summarise(rng = max(proportion_mean) - min(proportion_mean)) |>
#   arrange(desc(rng)) |>
#   slice_head(n = 6) |>
#   pull(cell_group)
# 
# raw <- counts_pt |>
#   group_by(sample) |>
#   mutate(prop = count / sum(count)) |>
#   ungroup() |>
#   filter(cell_group %in% top_cgs)
# 
# pred |>
#   filter(cell_group %in% top_cgs) |>
#   ggplot(aes(pseudotime, proportion_mean, colour = type, fill = type)) +
#   geom_ribbon(aes(ymin = proportion_lower, ymax = proportion_upper),
#               alpha = 0.15, linewidth = 0) +
#   geom_line(linewidth = 0.7) +
#   geom_point(data = raw, aes(y = prop), alpha = 0.5, size = 1.2) +
#   facet_wrap(~ cell_group, scales = "free_y", ncol = 3) +
#   theme_minimal(base_size = 10)

## ----test-smooth-endpoints, eval = run_smooth_vignette------------------------
# fit |>
#   sccomp_test_smooth(
#     smooth = "pseudotime",
#     from   = 1,
#     to     = 5
#   ) |>
#   arrange(c_FDR) |>
#   select(cell_group, from, to, c_effect, c_lower, c_upper, c_FDR)

## ----test-smooth-average, eval = run_smooth_vignette--------------------------
# fit |>
#   sccomp_test_smooth(
#     smooth     = "pseudotime",
#     from       = c(0, 2),   # early interval
#     to         = c(4, 6),   # late interval
#     comparison = "average"
#   ) |>
#   arrange(c_FDR) |>
#   select(cell_group, from, to, c_effect, c_lower, c_upper, c_FDR)

## ----fit-smooth-k3, message=FALSE, warning=FALSE, eval = run_smooth_vignette----
# fit_k3 <- counts_pt |>
#   sccomp_estimate(
#     formula_composition = ~ type + s(pseudotime, k = 3),
#     formula_variability = ~ 1,
#     sample              = "sample",
#     cell_group          = "cell_group",
#     abundance           = "count",
#     cores               = 1,
#     verbose             = FALSE,
#     max_sampling_iterations = 2000
#   )

## ----predict-grid-k3, message=FALSE, warning=FALSE, eval = run_smooth_vignette----
# pred_k3 <- fit_k3 |>
#   sccomp_predict(new_data = grid, number_of_draws = 200)

## ----plot-curves-k3, fig.width=8, fig.height=5, eval = run_smooth_vignette----
# pred_k3 |>
#   filter(cell_group %in% top_cgs) |>
#   ggplot(aes(pseudotime, proportion_mean, colour = type, fill = type)) +
#   geom_ribbon(aes(ymin = proportion_lower, ymax = proportion_upper),
#               alpha = 0.15, linewidth = 0) +
#   geom_line(linewidth = 0.7) +
#   geom_point(data = raw, aes(y = prop, colour = type), alpha = 0.5, size = 1.2) +
#   facet_wrap(~ cell_group, scales = "free_y", ncol = 3) +
#   theme_minimal(base_size = 10)

## ----basis-decomp, fig.width=8, fig.height=6, eval = run_smooth_vignette------
# smooth_meta    <- sccomp:::get_smooth_results(fit)
# smooth_specs   <- smooth_meta$smooth_specs
# smooth_re_objs <- smooth_meta$smooth_re_objs
# 
# grid_basis <- tibble(
#   pseudotime = seq(min(pt_map$pseudotime), max(pt_map$pseudotime),  length.out = 200),
#   type = levels(counts_obj$type)[1]  # any constant works; smooth ignores it
# )
# 
# basis <- sccomp:::predict_smooth_at_newdata(
#   smooth_specs[[1]], smooth_re_objs[[1]], grid_basis
# )
# # Plain s() has one penalised block — for multi-penalty smooths
# # (e.g. bs = "fs", t2(...)) `Xr_blocks` would contain several matrices.
# Xr_single <- basis$Xr_blocks[[1]]
# 
# basis_long <- bind_rows(
#   as.data.frame(basis$Xf) |>
#     setNames(sprintf("lin%d", seq_len(ncol(basis$Xf)))) |>
#     mutate(pseudotime = grid_basis$pseudotime) |>
#     pivot_longer(-pseudotime, names_to = "column", values_to = "value") |>
#     mutate(block = "Xf  (unpenalised, joins fixed effects)"),
#   as.data.frame(Xr_single) |>
#     setNames(sprintf("basis%d", seq_len(ncol(Xr_single)))) |>
#     mutate(pseudotime = grid_basis$pseudotime) |>
#     pivot_longer(-pseudotime, names_to = "column", values_to = "value") |>
#     mutate(block = "Xr  (penalised, joins one RE slot, shared sds)")
# )
# 
# ggplot(basis_long, aes(pseudotime, value, colour = column)) +
#   geom_line(linewidth = 0.7) +
#   facet_wrap(~ block, scales = "free_y", ncol = 1) +
#   labs(x = "pseudotime", y = "basis value",
#        title = "Decomposition of s(pseudotime, k = 5)",
#        subtitle = "Linear null-space (Xf) is a fixed effect; wiggly basis (Xr) is one RE slot") +
#   theme_minimal(base_size = 10)

## ----make-tissue, message=FALSE-----------------------------------------------
set.seed(11)
samples <- levels(counts_obj$sample)
covs <- tibble::tibble(
  sample     = samples,
  pseudotime = seq(0, 6, length.out = length(samples)) +
                 rnorm(length(samples), 0, 0.05),
  tissue     = factor(rep(c("blood", "lymph", "tumor"),
                          length.out = length(samples)))
)
counts_tissue <- counts_obj |> dplyr::left_join(covs, by = "sample")
counts_tissue |> dplyr::distinct(sample, type, pseudotime, tissue) |> head()

## ----fs-fit, message=FALSE, warning=FALSE, eval = run_smooth_vignette---------
# fit_fs <- counts_tissue |>
#   sccomp_estimate(
#     formula_composition = ~ s(pseudotime, tissue, bs = "fs", k = 5),
#     formula_variability = ~ 1,
#     sample              = "sample",
#     cell_group          = "cell_group",
#     abundance           = "count",
#     cores               = 1,
#     verbose             = FALSE,
#     max_sampling_iterations = 2000
#   )
# 
# # This one smooth occupies 3 RE slots (one per mgcv penalty block);
# # you can confirm via:
# attr(fit_fs, "model_input")$ncol_X_random_eff

## ----fs-predict, message=FALSE, warning=FALSE, eval = run_smooth_vignette-----
# grid_fs <- expand.grid(
#   pseudotime = seq(min(covs$pseudotime), max(covs$pseudotime),
#                    length.out = 60),
#   tissue     = levels(covs$tissue)
# ) |>
#   as_tibble() |>
#   mutate(sample = sprintf("fs_grid_%03d", row_number()))
# 
# pred_fs <- fit_fs |>
#   sccomp_predict(new_data = grid_fs, number_of_draws = 200)
# 
# head(pred_fs)

## ----fs-plot, fig.width=8.5, fig.height=5.5, eval = run_smooth_vignette-------
# top_cgs_fs <- pred_fs |>
#   group_by(cell_group) |>
#   summarise(rng = max(proportion_mean) - min(proportion_mean)) |>
#   arrange(desc(rng)) |>
#   slice_head(n = 6) |>
#   pull(cell_group)
# 
# raw_fs <- counts_tissue |>
#   group_by(sample) |>
#   mutate(prop = count / sum(count)) |>
#   ungroup() |>
#   filter(cell_group %in% top_cgs_fs)
# 
# pred_fs |>
#   filter(cell_group %in% top_cgs_fs) |>
#   ggplot(aes(pseudotime, proportion_mean,
#              colour = tissue, fill = tissue)) +
#   geom_ribbon(aes(ymin = proportion_lower, ymax = proportion_upper),
#               alpha = 0.15, linewidth = 0) +
#   geom_line(linewidth = 0.7) +
#   geom_point(data = raw_fs, aes(y = prop, colour = tissue),
#              alpha = 0.5, size = 1.2) +
#   facet_wrap(~ cell_group, scales = "free_y", ncol = 3) +
#   labs(x = "pseudotime", y = "proportion",
#        title = "Factor smooth s(pseudotime, tissue, bs = 'fs', k = 5)",
#        subtitle = "Per-tissue curves sharing one factor-smooth structure") +
#   theme_minimal(base_size = 10)

## ----multi-smooth-formula, eval=FALSE-----------------------------------------
# ~ type +
#   s(pseudotime, k = 5) +
#   s(age,        k = 4) +
#   (1 | donor)

## ----session-info, echo=FALSE-------------------------------------------------
sessionInfo()