# test-multiphase-parity.R — C binary parity and round-trip tests for multiphase
#
# Tests the R multiphase implementation against:
#   1. C HAZARD binary reference output (KUL CABG dataset)
#   2. Synthetic round-trip recovery
#
# The C reference model (hz.deadp.KUL) has ALL shapes fixed and only
# estimates 3 log(mu) parameters.  Our R implementation estimates shapes
# too, so parity tests focus on:
#   - Log-likelihood evaluation at C parameters (exact match)
#   - Full fit convergence with informed starting values
#   - Predict consistency

# ============================================================================
# Helpers
# ============================================================================

load_kul_csv <- function() {
  csv_path <- system.file("extdata", "cabgkul.csv", package = "TemporalHazard")
  if (!nzchar(csv_path) || !file.exists(csv_path)) {
    csv_path <- file.path("inst", "extdata", "cabgkul.csv")
  }
  if (!file.exists(csv_path)) return(NULL)
  utils::read.csv(csv_path)
}

load_synthetic_fixture <- function() {
  fixture_file <- system.file("fixtures", "mp_synthetic_3phase.rds",
                              package = "TemporalHazard")
  if (!nzchar(fixture_file) || !file.exists(fixture_file)) {
    fixture_file <- file.path("inst", "fixtures", "mp_synthetic_3phase.rds")
  }
  if (!file.exists(fixture_file)) return(NULL)
  readRDS(fixture_file)
}

# Starting theta on internal scale, matching SAS PARMS from hz.deadp.KUL.sas:
#   MUE=0.02, THALF=0.2, NU=1, M=1, MUC=0.0008, MUL=1E-9, TAU=1, GAMMA=3, ALPHA=1, ETA=1
# Internal layout: [log_mu, log_t_half, nu, m] per cdf/hazard phase,
#                  [log_mu] for constant,
#                  [log_mu, log_tau, gamma, alpha, eta] for g3 phase
kul_theta_start <- c(
  log(0.02),     # early.log_mu
  log(0.2),      # early.log_t_half
  1,             # early.nu
  1,             # early.m
  log(0.0008),   # constant.log_mu
  log(1e-9),     # late.log_mu
  log(1),        # late.log_tau
  3,             # late.gamma
  1,             # late.alpha
  1              # late.eta
)

kul_phases <- function() {
  list(
    early    = hzr_phase("cdf",  t_half = 0.2, nu = 1, m = 1),
    constant = hzr_phase("constant"),
    late     = hzr_phase("g3",   tau = 1, gamma = 3, alpha = 1, eta = 1)
  )
}


# ============================================================================
# C binary parity: log-likelihood at reference parameters
# ============================================================================

test_that("R log-likelihood at C reference parameters matches C output", {
  skip_on_cran()

  dat <- load_kul_csv()
  skip_if(is.null(dat), "KUL dataset not available")

  time   <- dat$int_dead
  status <- dat$dead

  phases <- kul_phases()
  covariate_counts <- c(early = 0L, constant = 0L, late = 0L)
  x_list <- list(early = NULL, constant = NULL, late = NULL)

  # C converged parameters on internal scale
  theta_c <- c(
    -3.77955,     # early.log_mu  (E0)
    log(0.2),     # early.log_t_half
    1,            # early.nu
    1,            # early.m
    -7.2258,      # constant.log_mu (C0)
    -16.6578,     # late.log_mu (L0)
    log(1),       # late.log_tau
    3,            # late.gamma
    1,            # late.alpha
    1             # late.eta
  )

  logl_r <- .hzr_logl_multiphase(
    theta = theta_c,
    time = time,
    status = status,
    phases = phases,
    covariate_counts = covariate_counts,
    x_list = x_list
  )

  expect_true(is.finite(logl_r))

  # C reference log-likelihood: -3740.52
  # R now uses G3 decomposition for the late phase, matching C.
  expect_equal(
    logl_r,
    -3740.52,
    tolerance = 0.01,
    label = "R logl at C params matches C reference"
  )
})


test_that("R cumulative hazard decomposition is consistent at C parameters", {
  skip_on_cran()

  dat <- load_kul_csv()
  skip_if(is.null(dat), "KUL dataset not available")

  phases <- kul_phases()
  covariate_counts <- c(early = 0L, constant = 0L, late = 0L)
  x_list <- list(early = NULL, constant = NULL, late = NULL)

  theta_c <- c(
    -3.77955, log(0.2), 1, 1,
    -7.2258,
    -16.6578, log(1), 3, 1, 1
  )

  t_grid <- c(0.1, 0.5, 1, 2, 5, 10, 20, 50, 100, 200)

  decomp <- .hzr_multiphase_cumhaz(
    time = t_grid,
    theta = theta_c,
    phases = phases,
    covariate_counts = covariate_counts,
    x_list = x_list,
    per_phase = TRUE
  )

  # All contributions should be non-negative
  expect_true(all(decomp$early >= 0))
  expect_true(all(decomp$constant >= 0))
  expect_true(all(decomp$late >= 0))

  # Total should equal sum of components
  expect_equal(
    decomp$total,
    decomp$early + decomp$constant + decomp$late,
    tolerance = 1e-12
  )

  # Survival = exp(-H) should be in [0, 1]
  surv <- exp(-decomp$total)
  expect_true(all(surv >= 0 & surv <= 1))

  # Early contribution should saturate (CDF bounded by 1 * mu_early)
  mu_early <- exp(-3.77955)
  expect_true(all(decomp$early <= mu_early + 1e-10))

  # Constant grows linearly with time
  mu_const <- exp(-7.2258)
  expect_equal(decomp$constant, mu_const * t_grid, tolerance = 1e-10)
})


test_that("Conservation of events at C reference parameters", {
  skip_on_cran()

  dat <- load_kul_csv()
  skip_if(is.null(dat), "KUL dataset not available")

  time   <- dat$int_dead
  status <- dat$dead

  phases <- kul_phases()
  covariate_counts <- c(early = 0L, constant = 0L, late = 0L)
  x_list <- list(early = NULL, constant = NULL, late = NULL)

  theta_c <- c(
    -3.77955, log(0.2), 1, 1,
    -7.2258,
    -16.6578, log(1), 3, 1, 1
  )

  cumhaz <- .hzr_multiphase_cumhaz(
    time = time,
    theta = theta_c,
    phases = phases,
    covariate_counts = covariate_counts,
    x_list = x_list
  )

  # C reference reports: "Number of events conserved = 544.9993"
  conserved <- sum(1 - exp(-cumhaz))
  expect_equal(conserved, 545, tolerance = 2)
})


# ============================================================================
# C binary parity: full model fit on KUL data
# ============================================================================

test_that("Full multiphase fit on KUL data converges", {
  skip_on_cran()

  dat <- load_kul_csv()
  skip_if(is.null(dat), "KUL dataset not available")

  # Provide starting theta from SAS PARMS so the optimizer begins in the

  # right basin.  Without this, the default mu_start = 0.01 for all phases
  # is far from the late-phase mu ≈ 5.8e-8 (log_mu ≈ -16.7).
  fit <- hazard(
    time   = dat$int_dead,
    status = dat$dead,
    dist   = "multiphase",
    theta  = kul_theta_start,
    phases = kul_phases(),
    fit = TRUE,
    control = list(n_starts = 10, maxit = 2000, reltol = 1e-10)
  )

  expect_s3_class(fit, "hazard")
  expect_true(is.finite(fit$fit$objective))

  # With 10 free parameters (shapes free), the optimizer can improve
  # beyond the shapes-fixed solution.  Should reach at least C's -3740.52.
  expect_true(fit$fit$objective >= -3745,
              label = paste("R logl", round(fit$fit$objective, 2),
                            ">= -3745"))

  # Phase names should be preserved
  expect_equal(names(fit$spec$phases), c("early", "constant", "late"))
  expect_equal(length(fit$fit$theta), 10)  # 4 + 1 + 5
})

# ============================================================================
# Shape-fixed fit: matches C binary workflow (only mu estimated)
# ============================================================================

test_that("KUL fit with fixed shapes: mechanism works correctly", {
  skip_on_cran()

  dat <- load_kul_csv()
  skip_if(is.null(dat), "KUL dataset not available")

  # C binary fixes all shapes and only estimates 3 mu parameters.
  # Replicate this by using fixed = "shapes" on each non-constant phase.
  phases_fixed <- list(
    early    = hzr_phase("cdf",  t_half = 0.2, nu = 1, m = 1,
                          fixed = "shapes"),
    constant = hzr_phase("constant"),
    late     = hzr_phase("g3",   tau = 1, gamma = 3, alpha = 1, eta = 1,
                          fixed = "shapes")
  )

  # --- Verify phases carry fixed attribute ---
  expect_equal(phases_fixed$early$fixed, c("t_half", "nu", "m"))
  expect_equal(phases_fixed$late$fixed, c("tau", "gamma", "alpha", "eta"))

  # Verify free_mask: 10 params total, 3 free (the 3 log_mu)
  cov_cnt <- c(early = 0L, constant = 0L, late = 0L)
  direct_mask <- .hzr_phase_free_mask(phases_fixed, cov_cnt)
  expect_equal(sum(direct_mask), 3L)
  expect_equal(direct_mask,
               c(TRUE, FALSE, FALSE, FALSE, TRUE, TRUE, FALSE, FALSE, FALSE, FALSE))

  # Starting values
  theta_start <- c(
    log(0.02),  log(0.2), 1, 1,
    log(0.0008),
    log(1e-7),  log(1),   3, 1, 1
  )

  fit <- hazard(
    time   = dat$int_dead,
    status = dat$dead,
    dist   = "multiphase",
    theta  = theta_start,
    phases = phases_fixed,
    fit    = TRUE,
    control = list(n_starts = 5, maxit = 2000, reltol = 1e-12)
  )

  expect_s3_class(fit, "hazard")
  expect_true(is.finite(fit$fit$objective))

  # Verify the optimizer used fixed-parameter masking
  expect_true(!is.null(fit$fit$fixed_mask),
              label = "fixed_mask should be set when shapes are fixed")

  # Shape-fixed fit with G3 late phase should match C's -3740.52
  expect_equal(fit$fit$objective, -3740.52, tolerance = 0.05,
               label = "R logl (shapes fixed) matches C reference")

  # Fixed shape parameters MUST stay at their starting values
  expect_equal(unname(fit$fit$theta[2]), log(0.2))   # early.log_t_half
  expect_equal(unname(fit$fit$theta[3]), 1)           # early.nu
  expect_equal(unname(fit$fit$theta[4]), 1)           # early.m
  expect_equal(unname(fit$fit$theta[7]), log(1))      # late.log_tau
  expect_equal(unname(fit$fit$theta[8]), 3)           # late.gamma
  expect_equal(unname(fit$fit$theta[9]), 1)           # late.alpha
  expect_equal(unname(fit$fit$theta[10]), 1)          # late.eta

  # Full theta should still be length 10 (shapes present but fixed)
  expect_equal(length(fit$fit$theta), 10)
})


test_that("KUL fit mu estimates are in neighborhood of C reference", {
  skip_on_cran()

  dat <- load_kul_csv()
  skip_if(is.null(dat), "KUL dataset not available")

  fit <- hazard(
    time = dat$int_dead,
    status = dat$dead,
    dist = "multiphase",
    theta = kul_theta_start,
    phases = kul_phases(),
    fit = TRUE,
    control = list(n_starts = 10, maxit = 2000, reltol = 1e-10)
  )

  theta <- fit$fit$theta

  # Internal layout: early=[log_mu, log_t_half, nu, m], constant=[log_mu],
  #                  late=[log_mu, log_tau, gamma, alpha, eta]
  log_mu_early <- theta[1]
  log_mu_const <- theta[5]
  log_mu_late  <- theta[6]

  # C reference: E0=-3.77955, C0=-7.2258, L0=-16.6578
  # With shapes free, mu may shift to compensate — use loose tolerance.
  # The early and constant phases are well-identified; late is harder.
  # G3 late phase has 4 shape params (tau, gamma, alpha, eta), creating
  # a wider mu-shape tradeoff surface than G1's 3 params, so late needs
  # extra tolerance.  Fixed-shape parity test validates exact logl match.
  expect_true(
    abs(log_mu_early - (-3.77955)) < 3,
    label = paste(
      "early log_mu", round(log_mu_early, 3),
      "within 3 of C ref -3.780"
    )
  )
  expect_true(
    abs(log_mu_const - (-7.2258)) < 3,
    label = paste(
      "constant log_mu", round(log_mu_const, 3),
      "within 3 of C ref -7.226"
    )
  )
  # Late log_mu has wide mu-shape tradeoff with 4 free G3 params.
  # Validate via log-likelihood rather than point estimate.
  # Free shapes should match or improve on C reference (which fixed shapes).
  expect_true(
    fit$fit$objective >= -3750,
    label = paste(
      "free-shape log-lik", round(fit$fit$objective, 2),
      "should be >= -3750 (C ref with fixed shapes: -3740.52)"
    )
  )
})

test_that("predict works on KUL multiphase fit", {
  skip_on_cran()

  dat <- load_kul_csv()
  skip_if(is.null(dat), "KUL dataset not available")

  fit <- hazard(
    time   = dat$int_dead,
    status = dat$dead,
    dist   = "multiphase",
    theta  = kul_theta_start,
    phases = kul_phases(),
    fit = TRUE,
    control = list(n_starts = 10, maxit = 2000, reltol = 1e-10)
  )

  # Predict at landmark times (months)
  nd <- data.frame(time = c(0.5, 1, 6, 12, 24, 60, 120, 200))
  surv <- predict(fit, newdata = nd, type = "survival")

  expect_equal(length(surv), 8)
  expect_true(all(surv >= 0 & surv <= 1))
  # Survival should be monotone decreasing
  expect_true(all(diff(surv) <= 0))

  # Decomposed cumulative hazard
  decomp <- predict(fit, newdata = nd, type = "cumulative_hazard",
                    decompose = TRUE)
  expect_true(is.data.frame(decomp))
  expect_true(all(c("time", "total", "early", "constant", "late") %in% names(decomp)))
  expect_equal(decomp$total, decomp$early + decomp$constant + decomp$late,
               tolerance = 1e-10)
})


# ============================================================================
# Synthetic golden fixture: round-trip recovery
# ============================================================================

test_that("Synthetic 3-phase fixture round-trip: parameters recovered", {
  skip_on_cran()

  fixture <- load_synthetic_fixture()
  skip_if(is.null(fixture), "Synthetic 3-phase fixture not available")

  # Skip if fixture was generated with divergent optimizer (sanity check)
  skip_if(is.null(fixture$fit), "Fixture has no fit results")
  skip_if(!is.finite(fixture$fit$logl), "Fixture fit did not converge")
  skip_if(
    any(abs(fixture$fit$theta) > 100),
    "Fixture theta unreasonable; regenerate fixtures (data-raw/golden_fixtures.R)"
  )

  # Refit using the same phase specs AND the fixture's converged theta as start
  refit <- hazard(
    time   = fixture$data$time,
    status = fixture$data$status,
    dist   = "multiphase",
    theta  = fixture$fit$theta,
    phases = list(
      early    = hzr_phase("cdf",
                           t_half = fixture$phases$early$t_half,
                           nu = fixture$phases$early$nu,
                           m = fixture$phases$early$m),
      constant = hzr_phase("constant"),
      late     = hzr_phase("hazard",
                           t_half = fixture$phases$late$t_half,
                           nu = fixture$phases$late$nu,
                           m = fixture$phases$late$m)
    ),
    fit = TRUE,
    # This round-trip check validates recovery from a stored converged theta.
    # Keep optimization deterministic and disable CoE reparameterization,
    # which can move theta along a near-flat ridge while preserving objective.
    control = list(n_starts = 1, conserve = FALSE, maxit = 1000, reltol = 1e-8)
  )

  expect_true(!is.null(refit$fit))

  # Parameter estimates should match fixture within tolerance
  expect_equal(refit$fit$theta, fixture$fit$theta, tolerance = 1e-3)

  # Log-likelihood should match
  expect_equal(refit$fit$objective, fixture$fit$logl, tolerance = 1e-2)
})


# ============================================================================
# Two-phase parity: early + constant (simpler case)
# ============================================================================

test_that("Two-phase (cdf + constant) round-trip on simulated data", {
  skip_on_cran()

  set.seed(123)
  n <- 300

  # True model: early CDF + constant
  mu_e <- 0.1
  t_half_e <- 1
  nu_e <- 1.5
  m_e <- 0
  mu_c <- 0.005

  t_grid <- seq(0.01, 30, length.out = 2000)
  d <- hzr_decompos(t_grid, t_half = t_half_e, nu = nu_e, m = m_e)
  H_true <- mu_e * d$G + mu_c * t_grid
  S_true <- exp(-H_true)

  u <- runif(n)
  time <- approx(S_true, t_grid, xout = u, rule = 2)$y
  cens_time <- runif(n, 0, quantile(time, 0.75))
  status <- as.integer(time <= cens_time)
  time <- pmin(time, cens_time)
  keep <- time > 0
  time <- time[keep]
  status <- status[keep]

  # Provide reasonable starting theta:
  #   early: [log(mu_e), log(t_half_e), nu_e, m_e] = [log(0.1), log(1), 1.5, 0]
  #   bg:    [log(mu_c)] = [log(0.005)]
  theta_start <- c(log(mu_e), log(t_half_e), nu_e, m_e, log(mu_c))

  fit <- hazard(
    time = time,
    status = status,
    dist = "multiphase",
    theta = theta_start,
    phases = list(
      early = hzr_phase("cdf", t_half = t_half_e, nu = nu_e, m = m_e),
      bg = hzr_phase("constant")
    ),
    fit = TRUE,
    control = list(n_starts = 3, maxit = 500)
  )

  expect_true(is.finite(fit$fit$objective))

  # Recovered mu_early should be in the right ballpark
  mu_early_hat <- exp(fit$fit$theta[1])
  expect_true(
    abs(log(mu_early_hat / mu_e)) < 1.5,
    label = paste(
      "early mu_hat", round(mu_early_hat, 4),
      "within 1.5 log-units of truth", mu_e
    )
  )

  # Predictions should be valid
  nd <- data.frame(time = c(0.5, 2, 10))
  surv <- predict(fit, newdata = nd, type = "survival")
  expect_true(all(surv >= 0 & surv <= 1))
  expect_true(all(diff(surv) <= 0))
})


# ============================================================================
# Standard error comparison (C reference, profile likelihood)
# ============================================================================

test_that("R standard errors at C params are consistent with C reference", {
  skip_on_cran()

  dat <- load_kul_csv()
  skip_if(is.null(dat), "KUL dataset not available")
  skip_if_not_installed("numDeriv")

  phases <- kul_phases()
  covariate_counts <- c(early = 0L, constant = 0L, late = 0L)
  x_list <- list(early = NULL, constant = NULL, late = NULL)

  # C reference converged theta (shapes fixed at exact values)
  theta_c <- c(-3.77955, log(0.2), 1, 1,
               -7.2258,
               -16.6578, log(1), 3, 1, 1)

  # Profile likelihood: only vary the 3 log_mu parameters (indices 1, 5, 6)
  # holding shapes fixed — matches what the C binary did.
  mu_idx <- c(1L, 5L, 6L)

  profile_logl <- function(mu_theta) {
    theta_full <- theta_c
    theta_full[mu_idx] <- mu_theta
    .hzr_logl_multiphase(
      theta = theta_full,
      time = dat$int_dead,
      status = dat$dead,
      phases = phases,
      covariate_counts = covariate_counts,
      x_list = x_list
    )
  }

  # Compute the early-phase SE via profile Hessian.  The late-phase
  # Hessian element is not numerically resolvable with finite differences
  # because mu_late ≈ 5.8e-8 contributes negligibly to the likelihood at
  # most observation times — the curvature signal is below double-precision
  # noise.  The C binary used analytic derivatives and had no difficulty.
  # TODO: add analytic Hessian; then test all 3 SEs.

  # Test the early-phase SE only (well-conditioned)
  early_logl <- function(log_mu_e) {
    theta_full <- theta_c
    theta_full[1] <- log_mu_e
    .hzr_logl_multiphase(
      theta = theta_full,
      time = dat$int_dead,
      status = dat$dead,
      phases = phases,
      covariate_counts = covariate_counts,
      x_list = x_list
    )
  }

  h <- 1e-4
  ll_p <- early_logl(theta_c[1] + h)
  ll_m <- early_logl(theta_c[1] - h)
  ll_0 <- early_logl(theta_c[1])
  d2 <- (ll_p + ll_m - 2 * ll_0) / h^2

  skip_if(
    !is.finite(d2) || d2 >= 0,
    "Early-phase Hessian not negative definite"
  )

  se_early <- sqrt(1 / (-d2))

  # C reference: SE(E0) = 0.09381214
  expect_equal(
    se_early,
    0.09381214,
    tolerance = 0.01,
    label = "Early phase SE matches C reference"
  )
})
