# This example demonstrates how to fit a proportional slowing
# progression model for repeated measures
# from first principles using the RTMB package.
# This example showcases the capabilities of RTMB
# for complicated nonlinear models with multivariate likelihoods,
# natural cubic splines, and sparse model matrices.

library(dplyr)
library(RTMB)

set.seed(0)
visit_times <- c(0, 12, 24, 36, 52, 64, 76)
n_patients <- 100
n_visits <- length(visit_times)

# For this example, we only use the pmrm package to simulate data.
# The actual modeling in this example calls RTMB directly.
data <- pmrm::pmrm_simulate_slowing_proportional(
  patients = n_patients,
  visit_times = visit_times,
  alpha = c(50, 44, 38, 32, 30, 28, 25),
  beta = c(0, 0.3), # 30% slowing due to treatment relative to placebo
  sigma = rep(3, n_visits),
  rho = rep(1, n_visits * (n_visits - 1) / 2),
  tau = 3
) |>
  mutate(
    arm = factor(
      arm,
      levels = c("arm_1", "arm_2"),
      labels = c("placebo", "treatment")
    )
  ) |>
  # Simulate a sparse model matrix for covariate adjustment:
  mutate(
    x1 = sample(c(0L, 1L), n(), replace = TRUE, prob = c(0.9, 0.1)),
    x2 = sample(c(0L, 1L), n(), replace = TRUE, prob = c(0.8, 0.2))
  )

constants <- list(
  y = data$y,
  spline_knots = visit_times,
  # RTMB integrates with sparse matrices from the Matrix package:
  X = Matrix::sparse.model.matrix(~ 0 + x1 + x2, data = data),
  arm = as.integer(factor(data$arm, levels = c("placebo", "treatment"))),
  time = data$t,
  n_visits = n_visits,
  n_patients = n_patients
)
constants$X_column_means <- Matrix::colMeans(constants$X)

parameters <- list(
  alpha = rep(40, n_visits), # Spline vertical knots
  beta = rep(0, 2), # Treatment effect parameters (slowing due to treatment)
  gamma = rep(0, 2), # Covariate adjustment fixed effect parameters
  phi = rep(0, n_visits), # Log of visit-specific standard deviation parameters
  rho = rep(0, n_visits * (n_visits - 1) / 2) # Latent correlation parameters
)

objective <- function(constants, parameters) {
  getAll(parameters, constants)
  y <- OBS(y)

  # Spline fitted to placebo respones:
  spline <- splinefun(x = spline_knots, y = alpha, method = "natural")

  # Covariate adjustment with sparse matrices (Matrix R package):
  adjustment <- (X %*% gamma) - sum(X_column_means * gamma)

  # Vector of fitted responses:
  mu <- spline((1 - beta[arm]) * time) + adjustment

  # Unstructured covariance among visits:
  sigma <- diag(exp(phi))
  Sigma <- sigma %*% unstructured(n_visits)$corr(rho) %*% sigma

  # Multivariate normal likelihood (on the minus log scale):
  target <- 0
  for (patient in seq_len(n_patients)) {
    visits <- seq(to = patient * n_visits, length.out = n_visits)
    target <- target - RTMB::dmvnorm(y[visits], mu[visits], Sigma, log = TRUE)
  }
  target
}

# Register the objective function with RTMB:
model <- MakeADFun(
  func = function(parameters) objective(constants, parameters),
  parameters = parameters,
  map = list(beta = factor(c(NA, 1))), # Set beta[1] constant (at 0).
  silent = TRUE
)

# Fit the model using an optimizer in R:
optimization <- nlminb(model$par, model$fn, model$gr)

# Check convergence:
print(optimization$convergence) # 0 indicates successful convergence

# Get parameter estimates:
report <- sdreport(model)
print(as.list(report, "Est")) # parameter estimates
print(as.list(report, "Std")) # standard errors