# Conditional logit (discrete-choice REM) -------------------------------------
#
# Ordinal (clogit):
#   log-likelihood = Σ_t  β'x_{i_t j_t}  -  log Σ_{(i,j) ∈ R_t} exp(β'x_{ij})
#   which is McFadden's conditional logit.
#
# Temporal (interval):
#   log-likelihood = Σ_t [ β'x_{i_t j_t}
#                         - log Σ_{(i,j)∈R_t} exp(β'x_{ij})          (choice term)
#                         - Δ_t × Σ_{(i,j)∈R_t} exp(β'x_{ij}) ]     (waiting-time penalty)
#   where Δ_t = t_t − t_{t−1} is the inter-event gap.
#   This is the full Hawkes/REM temporal likelihood; the choice term alone is
#   the ordinal likelihood.  Both share the same risk-set statistics.
#
# Both are maximised via optim(method="BFGS") with analytic gradients.

.rem_clogit_negll <- function(beta, X, strata) {
  eta    <- drop(X %*% beta)               # linear predictor, length = nrow(X)
  # subtract per-stratum max for numerical stability
  s_max  <- tapply(eta, strata, max)[strata]
  eta_c  <- eta - s_max
  exp_e  <- exp(eta_c)
  log_denom <- log(tapply(exp_e, strata, sum))[strata]
  lp     <- eta_c - log_denom              # log P(choose this row | stratum)
  ll     <- sum(lp[X[, 1L] == X[, 1L]])   # dummy; will use occurred mask below
  ll
}

# Fit conditional logit REM via BFGS.
# rs      : risk-set data frame with column 'occurred' and covariate columns
# formula : one-sided ~ covariate formula (no Surv(), no strata())
# Returns a list mimicking enough of coxph to be useful: coef, vcov, loglik
.fit_clogit <- function(rs, formula) {
  covars <- all.vars(formula)
  X      <- as.matrix(rs[, covars, drop = FALSE])
  y      <- rs$occurred
  g      <- rs$time2   # stratum = event time (unique per event)

  # negative log-likelihood
  negll <- function(beta) {
    eta   <- drop(X %*% beta)
    s_max <- tapply(eta, g, max)[as.character(g)]
    ec    <- exp(eta - s_max)
    ld    <- log(tapply(ec, g, sum))[as.character(g)]
    -sum((eta - s_max - ld)[y == 1L])
  }

  # analytic gradient
  grad <- function(beta) {
    eta   <- drop(X %*% beta)
    s_max <- tapply(eta, g, max)[as.character(g)]
    ec    <- exp(eta - s_max)
    denom <- tapply(ec, g, sum)[as.character(g)]
    p     <- ec / denom                   # P(row | stratum)
    # gradient = Σ_t [ x_{chosen_t} - E[x | stratum t] ]
    # = X'(y - p)  (where y is the 0/1 occurred indicator)
    -drop(t(X) %*% (y - p))
  }

  beta0 <- rep(0, ncol(X))
  fit   <- stats::optim(beta0, negll, grad, method = "BFGS",
                        control = list(maxit = 500L),
                        hessian = TRUE)

  coef <- fit$par
  names(coef) <- covars

  # vcov from observed Fisher information
  vcov <- tryCatch(solve(fit$hessian), error = function(e) {
    matrix(NA_real_, ncol(X), ncol(X),
           dimnames = list(covars, covars))
  })
  dimnames(vcov) <- list(covars, covars)

  structure(
    list(coef   = coef,
         vcov   = vcov,
         loglik = c(negll(rep(0, ncol(X))), -fit$value),
         converged = fit$convergence == 0L,
         method = "clogit"),
    class = "rem_clogit"
  )
}

# Fit temporal REM via BFGS.
#
# The temporal (interval-censored) log-likelihood is:
#   LL = Σ_{t=2}^{m} [ β'x_{focal,t}  −  Δ_t × Σ_{ij} exp(β'x_{ij,t}) ]
#        − Δ_{m+1} × Σ_{ij} exp(β'x_{ij,m})       ← right-censoring at end
#
# where Δ_t = t_t − t_{t−1} and the outer sum is over all n*(n-1) dyads.
# Unlike the ordinal likelihood there is no log-normaliser; the rate of the
# focal dyad enters directly, not as a conditional probability.
#
# With risk-set sampling the population total rate is estimated as:
#   Λ_t ≈ (n_dyads / ns_t) × Σ_{ij ∈ sample_t} exp(β'x_{ij,t})
# This is unbiased when exp(η) is roughly uniform across dyads; it degrades
# for very skewed rate distributions (a known limitation of rate-model sampling).
#
# rs      : risk-set data frame (remstats_generate output + add_survival_cols)
# formula : one-sided ~ covariate formula
# deltas  : numeric vector length m, delta[t] = t_t − t_{t−1}, delta[1] = t_1
# n_dyads : n*(n-1); used to scale the sampled total rate back to population
.fit_temporal <- function(rs, formula, deltas, n_dyads) {
  covars <- all.vars(formula)
  X      <- as.matrix(rs[, covars, drop = FALSE])
  y      <- rs$occurred
  g      <- rs$time2   # stratum id = event time (unique per event, events 2..m)

  strata_ord <- sort(unique(g))               # length = m - 1
  # deltas for strata 2..m are deltas[2..m]; index via position in strata_ord
  # strata_ord[k] is the k-th unique event time; its gap is deltas[k+1]
  # (deltas[1] = t_1, not used in the sum — event 1 has no risk set in pe)
  delta_u   <- deltas[seq(2L, length(deltas))]   # length = m - 1
  delta_map <- delta_u[match(g, strata_ord)]      # per risk-set row

  # sampling scale: n_dyads / ns_t per stratum
  ns_tab    <- table(g)
  scale_u   <- n_dyads / as.numeric(ns_tab)       # length = m - 1
  scale_map <- n_dyads / as.numeric(ns_tab[as.character(g)])

  negll <- function(beta) {
    eta <- drop(X %*% beta)
    # focal log-rates: Σ_t β'x_{focal,t}
    log_rate_focal <- sum(eta[y == 1L])
    # total rate per stratum: Λ_t ≈ scale_u * Σ_sample exp(η)
    # use log-sum-exp for stability
    s_max  <- tapply(eta, g, max)
    ec     <- exp(eta - s_max[as.character(g)])
    sum_ec <- as.numeric(tapply(ec, g, sum))       # length = m - 1
    Lambda <- scale_u * sum_ec * exp(s_max)
    # negative log-likelihood
    -(log_rate_focal - sum(delta_u * Lambda))
  }

  grad <- function(beta) {
    eta <- drop(X %*% beta)
    # focal contribution: -x_{focal,t} per event
    g_focal <- -colSums(X[y == 1L, , drop = FALSE])
    # penalty gradient: Σ_t Δ_t × scale_t × Σ_{ij∈sample_t} exp(η_ij) × x_ij
    ec_abs  <- exp(eta)
    g_pen   <- drop(t(X) %*% (delta_map * scale_map * ec_abs))
    g_focal + g_pen
  }

  beta0 <- rep(0, ncol(X))
  fit   <- stats::optim(beta0, negll, grad, method = "BFGS",
                        control = list(maxit = 500L),
                        hessian = TRUE)

  coef <- fit$par
  names(coef) <- covars

  vcov <- tryCatch(solve(fit$hessian), error = function(e) {
    matrix(NA_real_, ncol(X), ncol(X), dimnames = list(covars, covars))
  })
  dimnames(vcov) <- list(covars, covars)

  structure(
    list(coef      = coef,
         vcov      = vcov,
         loglik    = c(-negll(rep(0, ncol(X))), -fit$value),
         converged = fit$convergence == 0L,
         method    = "temporal"),
    class = "rem_clogit"
  )
}

#' @export
coef.rem_clogit <- function(object, ...) object$coef

#' @export
print.rem_clogit <- function(x, ...) {
  se  <- sqrt(diag(x$vcov))
  z   <- x$coef / se
  tab <- data.frame(coef = x$coef, `se(coef)` = se, z = z,
                    `Pr(>|z|)` = 2 * stats::pnorm(-abs(z)),
                    check.names = FALSE)
  label <- if (identical(x$method, "temporal")) "REM temporal" else "REM conditional logit"
  cat(label, "\n\n")
  print(round(tab, 4L))
  if (!x$converged) cat("\n[Warning: optimizer did not converge]\n")
  invisible(x)
}

# covariate formula from a Surv formula (strips Surv() and strata())
.surv_to_covar_formula <- function(formula) {
  tt  <- terms(formula, specials = c("strata", "survival::strata"))
  labs <- attr(tt, "term.labels")
  # drop any term whose label contains "strata("
  labs <- labs[!grepl("strata\\(", labs)]
  stats::as.formula(paste("~", paste(labs, collapse = " + ")))
}

#' Compute sufficient-statistic means for a sequence without fitting a REM
#'
#' Builds covariate values only for the realized events in `edgelist`
#' (n_nonevents = 0L) and returns their column means.  Much faster than
#' [fit_rem()] since no Cox model is fitted and the risk set contains exactly
#' one row per event.  Use to measure convergence in sufficient-statistic space
#' rather than coefficient space.
#'
#' @param edgelist    Prepared edgelist (output of [prepare_edgelist()]).
#' @param rem_formula Survival formula — used to identify covariate columns.
#' @param base_riskset Precomputed base risk set; constructed from `edgelist`
#'   if `NULL`.
#' @param covariates Optional exogenous covariate list from [rem_covariates()].
#'
#' @return Named numeric vector: mean of each REM covariate across events.
#' @examples
#' data(edgelist)
#' el <- prepare_edgelist(edgelist)
#' compute_seq_stat_means(
#'   el[1:100, ],
#'   survival::Surv(time1, occurred) ~ Repetition + RecRecencyLog +
#'     survival::strata(time2)
#' )
#' @export
compute_seq_stat_means <- function(edgelist, rem_formula, base_riskset = NULL,
                                    covariates = NULL) {
  all_nodes    <- get_nodes(edgelist, NULL)
  base_riskset <- get_base_riskset(base_riskset, all_nodes)

  pe <- remstats_generate(edgelist,
                           n            = length(all_nodes),
                           n_nonevents  = 0L,
                           base_riskset = base_riskset,
                           covariates   = covariates)
  pe <- add_survival_cols(pe)

  ss <- make_suffstats(rem_formula, pe, drop_strata = TRUE)
  colMeans(ss, na.rm = TRUE)
}

#' Fit a Relational Event Model to an edgelist
#'
#' Builds risk-set statistics for each event in the sequence, adds survival
#' model timing columns (`time1`, `time2`), and fits a Cox proportional hazards
#' model. All covariate computation is handled internally.
#'
#' @param edgelist Data frame with columns `time`, `sender`, `receiver`,
#'   `eveid`. Use [prepare_edgelist()] to create this from a raw matrix.
#' @param formula A `survival::Surv` formula. Must reference `time1` and
#'   `time2` (via `strata()`).
#'
#'   **Available covariates**
#'
#'   \tabular{llll}{
#'     \strong{Name} \tab \strong{Definition} \tab \strong{relevent analogue} \tab \strong{Notes} \cr
#'     \emph{Dyad activity} \tab \tab \tab \cr
#'     `Repetition` \tab dyad_occ / sender_outDeg \tab `FrPSndSnd` \tab Same quantity; llrem denom = n-1 (fixed), relevent = event index \cr
#'     `SndInitiation` \tab I(dyad i→j has ever fired) \tab — \tab llrem only; used as partition variable in [fit_rem_sep()] \cr
#'     `DyadRecencyLog` \tab -log1p(time since last i→j) \tab `RSndSnd` \tab Same direction; relevent uses recency rank, not log-time \cr
#'     `WeightedRep` \tab w_ij / max(w_ij in risk set) \tab — \tab llrem only; same-dyad count normalised by per-sender max \cr
#'     \emph{Reciprocity} \tab \tab \tab \cr
#'     `RRecInitiation` \tab I(reverse dyad j→i has ever fired) \tab `PSAB-BA` (one-step) \tab relevent PSAB-BA resets after each non-BA event; RRecInitiation is permanent \cr
#'     `RecRecencyLog` \tab -log1p(time since last j→i) \tab `RRecSnd` \tab Both measure reverse-dyad recency; relevent uses rank, llrem uses log-time \cr
#'     `RecRecencyRank` \tab 1 / rank(j in i's incoming recency list) \tab `RRecSnd` \tab Exact relevent analogue; rank 1 = most recent sender to i \cr
#'     `WeightedRec` \tab min(w_ij, w_ji) / max(w_ij, w_ji) \tab — \tab Squartini et al. (2013) weighted reciprocity; 0 if both zero \cr
#'     `FrRecSnd` \tab w_ji / sender_inDeg \tab `FrRecSnd` \tab Fraction of sender's in-events that are from receiver \cr
#'     \emph{Degree} \tab \tab \tab \cr
#'     `OutDegSnd` \tab sender out-degree / (n-1) \tab `NODSnd` \tab relevent divides by running event count \cr
#'     `OutDegRec` \tab receiver out-degree / (n-1) \tab `NODRec` \tab same normalisation difference \cr
#'     `InDegSnd` \tab sender in-degree / (n-1) \tab `NIDSnd` \tab same normalisation difference \cr
#'     `InDegRec` \tab receiver in-degree / (n-1) \tab `NIDRec` \tab same normalisation difference \cr
#'     `TotDegSnd` \tab (OutDegSnd + InDegSnd) / 2 \tab `NTDegSnd` \tab both are total-degree averages \cr
#'     `TotDegRec` \tab (OutDegRec + InDegRec) / 2 \tab `NTDegRec` \tab both are total-degree averages \cr
#'     \emph{Triadic closure} \tab \tab \tab \cr
#'     `OTP` \tab out-two-path count \tab `OTPSnd` \tab raw count; use `NOTPSnd` for normalised \cr
#'     `ITP` \tab in-two-path count \tab `ITPSnd` \tab raw count; use `NITPSnd` for normalised \cr
#'     `OSP` \tab out-shared-partner count \tab `OSPSnd` \tab raw count; use `NOSPSnd` for normalised \cr
#'     `ISP` \tab in-shared-partner count \tab `ISPSnd` \tab raw count; use `NISPSnd` for normalised \cr
#'     `NOTPSnd` \tab OTP / max(OTP in risk set) \tab — \tab llrem only (normalised OTP) \cr
#'     `NITPSnd` \tab ITP / max(ITP in risk set) \tab — \tab llrem only (normalised ITP) \cr
#'     `NOSPSnd` \tab OSP / max(OSP in risk set) \tab — \tab llrem only (normalised OSP) \cr
#'     `NISPSnd` \tab ISP / max(ISP in risk set) \tab — \tab llrem only (normalised ISP) \cr
#'     `NTriadic` \tab (OTP + ITP + OSP + ISP) / max(OTP + ITP + OSP + ISP) \tab — \tab llrem only; omnibus triadic score; ranges 0--1 \cr
#'     \emph{Imitation} \tab \tab \tab \cr
#'     `Imitation` \tab (recInDg - w_ij) / max per sender \tab — \tab Times j targeted by actors other than i, normalised; ranges 0--1 \cr
#'     `PreImitation` \tab recInDg - w_ij (raw, unnormalised) \tab — \tab Raw count underlying Imitation \cr
#'     \emph{P-shifts} \tab \tab \tab \cr
#'     `IOviolation` \tab PSAB_XA + PSAB_XB + PSAB_XY (any turn-usurp) \tab `PSAB-XA`/`PSAB-XB`/`PSAB-XY` (separate) \tab llrem collapses the three usurping types \cr
#'     `PSAB_BA` \tab I(last event was j→i, this is i→j) \tab `PSAB-BA` \tab exact match \cr
#'     `PSAB_BY` \tab I(last event was j→i, this is i→y, y!=j) \tab `PSAB-BY` \tab exact match \cr
#'     `PSAB_AY` \tab I(last event was i→j, this is i→y, y!=j) \tab `PSAB-AY` \tab exact match \cr
#'     `PSAB_XA` \tab I(last event was x→i, this is j→i) \tab `PSAB-XA` \tab exact match \cr
#'     `PSAB_XB` \tab I(last event was x→j, this is i→j) \tab `PSAB-XB` \tab exact match \cr
#'     `PSAB_XY` \tab I(last event was x→y, this is i→j; x!=i, y!=j) \tab `PSAB-XY` \tab exact match \cr
#'   }
#' @param n Number of nodes. Inferred from `edgelist` if `NULL`.
#' @param n_nonevents Sampled non-event risk-set size per event. `NULL` uses
#'   the full risk set (slow for large networks; 200 is a sensible default).
#' @param base_riskset Precomputed base risk-set data frame. Constructed from
#'   `n` if `NULL`. Pass a cached object when calling `fit_rem()` and
#'   [run_llm_rem()] in the same session to avoid redundant work.
#' @param cl Optional PSOCK cluster from `parallel::makeCluster()` for
#'   parallelised risk-set computation.
#' @param method Estimation method: `"cox"` (default) fits a Cox proportional
#'   hazards model via [survival::coxph()] with one stratum per event.
#'   `"clogit"` fits a conditional logit (discrete-choice/ordinal) model via
#'   BFGS; substantially faster than `"cox"` for large sequences and
#'   asymptotically equivalent under the ordinal likelihood.
#'   `"temporal"` fits the full interval-censored REM likelihood, which adds a
#'   waiting-time penalty `Delta_t * Sigma exp(eta)` to each event's contribution.
#'   Requires that `edgelist$time` contains real inter-event gaps, not just
#'   ranks; coefficients are no longer scale-free.
#'
#'   Available covariates: see `@param formula` above.
#'
#'   Note on naming vs **relevent**: `OutDegSnd`/`OutDegRec`/`InDegSnd`/`InDegRec`
#'   scale by `(n-1)` (fixed maximum degree), unlike relevent's `NODSnd`/`NIDRec`
#'   which scale by the running event count.  `DyadRecencyLog` is
#'   `-log1p(clock-time of last i→j event)`; relevent's `RSndSnd` is a recency
#'   rank -- both measure same-dyad recency but are not numerically comparable.
#'   `RecRecencyLog` (`-log1p(time since last j→i)`) and `RecRecencyRank`
#'   (`1/rank` of j in i's incoming recency list) are both analogues of
#'   relevent's `RRecSnd`; `RecRecencyRank` is the exact functional equivalent.
#'
#' @param covariates Optional exogenous covariate list created by
#'   [rem_covariates()]. `NULL` uses only the built-in endogenous statistics.
#' @return For `method = "cox"`: a `coxph` object.  For `method = "clogit"` or
#'   `"temporal"`: a `rem_clogit` object with elements `coef`, `vcov`,
#'   `loglik`, and `converged`.
#' @examples
#' data(edgelist)
#' el <- prepare_edgelist(edgelist)
#' fit <- fit_rem(
#'   el[1:100, ],
#'   survival::Surv(time1, occurred) ~ Repetition + RecRecencyLog +
#'     survival::strata(time2),
#'   n_nonevents = 20L, method = "clogit"
#' )
#' print(fit)
#' @export
fit_rem <- function(edgelist,
                    formula,
                    n            = NULL,
                    n_nonevents  = 200L,
                    base_riskset = NULL,
                    cl           = NULL,
                    method       = c("cox", "clogit", "temporal"),
                    covariates   = NULL) {
  method       <- match.arg(method)
  all_nodes    <- get_nodes(edgelist, n)
  base_riskset <- get_base_riskset(base_riskset, all_nodes)
  pe           <- remstats_generate(edgelist,
                                     n            = length(all_nodes),
                                     n_nonevents  = n_nonevents,
                                     cl           = cl,
                                     base_riskset = base_riskset,
                                     covariates   = covariates)
  pe <- add_survival_cols(pe)

  if (method %in% c("clogit", "temporal")) {
    covar_formula <- .surv_to_covar_formula(formula)
    if (method == "clogit")
      return(.fit_clogit(pe, covar_formula))
    # inter-event gaps in the same order as sort(unique(pe$time2))
    times   <- sort(unique(edgelist$time))
    deltas  <- c(times[1], diff(times))   # delta[1] = t_1 (first gap from 0)
    n_dyads <- length(all_nodes) * (length(all_nodes) - 1L)
    return(.fit_temporal(pe, covar_formula, deltas, n_dyads))
  }
  survival::coxph(formula, data = pe)
}

#' Generate an LLM event sequence and fit a REM
#'
#' At each time step *t* (starting from 2), `query_fn` is called with the
#' growing LLM-generated history, the empirical event at *t*, and the full
#' node set. The returned sender/receiver pair is appended to the sequence.
#' The sequence is seeded with the first empirical event so the LLM starts
#' with at least one data point. After generation, risk-set statistics are
#' built and a REM is fitted.
#'
#' @param edgelist Prepared edgelist (output of [prepare_edgelist()]).
#' @param formula Survival formula passed to [fit_rem()].
#' @param query_fn A closure produced by [make_query_roleplay_sender()],
#'   [make_query_predict()], or [make_query_roleplay_random()], or any function
#'   with the signature:
#'   `f(past_events, empirical_event, all_nodes) -> c(sender, receiver)`.
#' @param n_events Number of events to generate. `NULL` uses the full edgelist.
#' @param n Number of nodes. Inferred if `NULL`.
#' @param n_nonevents Sampled risk-set size passed to [fit_rem()].
#' @param base_riskset Precomputed base risk set.
#' @param cl Optional PSOCK cluster for parallelised risk-set computation.
#' @param seed RNG seed for random fallbacks and random-node sampling.
#' @param checkpoint_path File path for incremental saves every 10 events.
#'   `NULL` disables checkpointing.
#' @param use_empirical_history If `FALSE` (default), the LLM is shown its own
#'   previously generated events as history, so any divergence from reality
#'   compounds over time. If `TRUE`, the LLM is shown the empirical event
#'   history at each step, isolating the effect of the generation strategy from
#'   the effect of accumulated sequence drift.
#' @param rate_limit_secs Seconds to pause between LLM calls.
#' @param fit If `TRUE` (default), fits a REM on the generated sequence before
#'   returning. Set to `FALSE` to skip model fitting and return only the
#'   sequence and raw responses — useful when running many sequences and
#'   deferring all model fitting to a separate step.
#' @param method Passed to [fit_rem()]: `"cox"` (default), `"clogit"`, or
#'   `"temporal"`.
#' @param covariates Optional exogenous covariate list passed to [fit_rem()].
#'
#' @return A named list with elements:
#'   - `sequence`: generated event sequence as a data frame.
#'   - `responses`: raw LLM response text and parsed sender/receiver per event.
#'   - `rem`: fitted model (only present when `fit = TRUE`).
#' @examples
#' \dontrun{
#' data(edgelist)
#' el  <- prepare_edgelist(edgelist)
#' cfg <- rem_cfg(
#'   provider    = llm_provider_openai(),
#'   context     = "These are militarized conflict events between countries.",
#'   action_verb = "attacked",
#'   actor_noun  = "Country"
#' )
#' res <- run_llm_rem(
#'   el[1:50, ],
#'   survival::Surv(time1, occurred) ~ Repetition + RecRecencyLog +
#'     survival::strata(time2),
#'   query_fn    = make_query_roleplay_sender(cfg),
#'   n_nonevents = 50L, method = "clogit"
#' )
#' print(res$rem)
#' }
#' @export
run_llm_rem <- function(edgelist,
                         formula,
                         query_fn,
                         n_events              = NULL,
                         n                     = NULL,
                         n_nonevents           = 200L,
                         base_riskset          = NULL,
                         cl                    = NULL,
                         seed                  = 8214L,
                         checkpoint_path       = NULL,
                         use_empirical_history = FALSE,
                         rate_limit_secs       = 0.3,
                         fit                   = TRUE,
                         method                = c("cox", "clogit", "temporal"),
                         covariates            = NULL) {
  method <- match.arg(method)
  all_nodes    <- get_nodes(edgelist, n)
  base_riskset <- get_base_riskset(base_riskset, all_nodes)

  n_ev <- if (is.null(n_events)) nrow(edgelist) else min(n_events, nrow(edgelist))
  emp  <- edgelist[seq_len(n_ev), ]

  set.seed(seed)
  llm_rows <- list()
  start_t  <- 2L

  if (!is.null(checkpoint_path) && file.exists(checkpoint_path)) {
    env <- new.env(parent = emptyenv())
    load(checkpoint_path, envir = env)
    llm_rows <- env$llm_rows
    start_t  <- length(llm_rows) + 2L
    message("[llrem] Resuming from checkpoint at event ", start_t)
  }

  llm_seq <- emp[1L, c("time", "sender", "receiver", "eveid")]

  if (start_t > n_ev) {
    message("[llrem] Checkpoint is complete; skipping generation.")
  } else {
    pb <- utils::txtProgressBar(min = start_t, max = n_ev, style = 3)

    for (t in start_t:n_ev) {
      history <- if (use_empirical_history) emp[seq_len(t - 1L), ] else llm_seq
      result <- query_fn(
        past_events     = history,
        empirical_event = emp[t, ],
        all_nodes       = all_nodes
      )
      new_row <- data.frame(
        time     = emp$time[t],
        sender   = as.integer(result[["sender"]]),
        receiver = as.integer(result[["receiver"]]),
        eveid    = paste0(result[["sender"]], "_", result[["receiver"]]),
        stringsAsFactors = FALSE
      )
      llm_rows[[t - 1L]] <- new_row
      llm_seq <- rbind(llm_seq, new_row[, c("time", "sender", "receiver", "eveid")])

      if (!is.null(checkpoint_path) && (t %% 10L == 0L || t == n_ev))
        save(llm_rows, file = checkpoint_path)

      utils::setTxtProgressBar(pb, t)
      Sys.sleep(rate_limit_secs)
    }
    close(pb)
  }

  sequence <- rbind(emp[1L, c("time", "sender", "receiver", "eveid")],
                    do.call(rbind, llm_rows))
  row.names(sequence) <- NULL

  out <- list(sequence  = sequence,
              responses = do.call(rbind, llm_rows))

  if (fit) {
    message("\n[llrem] Fitting REM on generated sequence...")
    out$rem <- fit_rem(sequence, formula,
                       n = length(all_nodes), n_nonevents = n_nonevents,
                       base_riskset = base_riskset, cl = cl, method = method,
                       covariates = covariates)
  }

  out
}

# ---------------------------------------------------------------------------
# Separable REM
# ---------------------------------------------------------------------------

#' Fit a Separable Relational Event Model
#'
#' Partitions the risk set by prior contact status (`SndInitiation`) and fits
#' two independent sub-models: one for **initiation** events (sender has never
#' previously contacted receiver; `SndInitiation == 0`) and one for
#' **continuation** events (sender has contacted receiver before;
#' `SndInitiation == 1`).  The approach is analogous to separable TERGMs
#' (Krivitsky and Handcock, 2014): because the two subsets are disjoint, their
#' parameters are conditionally independent and can be estimated separately via
#' standard Cox partial likelihood on each partition.
#'
#' The combined log-likelihood is `loglik_init + loglik_cont`, which is directly
#' comparable to a pooled five-parameter model with the same degrees of freedom
#' (see paper for details).
#'
#' @param edgelist Data frame with columns `time`, `sender`, `receiver`,
#'   `eveid`. Use [prepare_edgelist()] to create this from a raw matrix.
#' @param formula_init Survival formula for the initiation sub-model
#'   (events where `SndInitiation == 0`).  Typical terms: `RRecInitiation`,
#'   `NOTPSnd`, `NITPSnd`.  Must include `survival::strata(time2)`.
#' @param formula_cont Survival formula for the continuation sub-model
#'   (events where `SndInitiation == 1`).  Typical terms: `Repetition`,
#'   `RecRecencyLog`.  Must include `survival::strata(time2)`.
#' @param n Number of nodes. Inferred from `edgelist` if `NULL`.
#' @param n_nonevents Sampled non-event risk-set size per event. `NULL` uses
#'   the full risk set. Passed to [fit_rem()] internally.
#' @param base_riskset Precomputed base risk-set data frame.
#' @param cl Optional PSOCK cluster for parallelised risk-set computation.
#' @param method Estimation method: `"cox"` (default), `"clogit"`, or
#'   `"temporal"`. Applied independently to each sub-model.
#' @param covariates Optional exogenous covariate list from [rem_covariates()].
#' @param min_events Minimum number of focal events required in each partition
#'   to attempt fitting. Partitions below this threshold return `NULL` for
#'   that sub-model with a warning. Default 2.
#'
#' @return A `rem_sep` object (a named list) with elements:
#'   \describe{
#'     \item{`init`}{Fitted initiation sub-model (same class as [fit_rem()]
#'       output for the chosen `method`), or `NULL` if too few events.}
#'     \item{`cont`}{Fitted continuation sub-model, or `NULL`.}
#'     \item{`loglik`}{Named numeric vector: `init`, `cont`, `sep`
#'       (= init + cont), and `null` (= init_null + cont_null).}
#'     \item{`n_init`}{Number of initiation events in the sequence.}
#'     \item{`n_cont`}{Number of continuation events.}
#'     \item{`init_fraction`}{Proportion of events that are initiations.}
#'     \item{`method`}{The estimation method used.}
#'   }
#' @examples
#' data(edgelist)
#' el <- prepare_edgelist(edgelist)
#' fit <- fit_rem_sep(
#'   el[1:200, ],
#'   formula_init = survival::Surv(time1, occurred) ~
#'     RRecInitiation + survival::strata(time2),
#'   formula_cont = survival::Surv(time1, occurred) ~
#'     Repetition + RecRecencyLog + survival::strata(time2),
#'   n_nonevents = 20L, method = "clogit"
#' )
#' print(fit)
#' @export
fit_rem_sep <- function(edgelist,
                        formula_init,
                        formula_cont,
                        n            = NULL,
                        n_nonevents  = 200L,
                        base_riskset = NULL,
                        cl           = NULL,
                        method       = c("cox", "clogit", "temporal"),
                        covariates   = NULL,
                        min_events   = 2L) {
  method       <- match.arg(method)
  all_nodes    <- get_nodes(edgelist, n)
  base_riskset <- get_base_riskset(base_riskset, all_nodes)

  pe <- remstats_generate(edgelist,
                           n            = length(all_nodes),
                           n_nonevents  = n_nonevents,
                           cl           = cl,
                           base_riskset = base_riskset,
                           covariates   = covariates)
  pe <- add_survival_cols(pe)

  # Partition ----------------------------------------------------------------
  # A stratum (event time) contains exactly one focal row; that row belongs
  # entirely to one partition.  Strata that fall in the other partition must
  # be dropped so each sub-model's partial likelihood is well-formed.
  pe_init <- pe[pe$SndInitiation == 0, ]
  pe_cont <- pe[pe$SndInitiation == 1, ]

  # Keep only strata that contain at least one focal event in the partition
  keep_strata <- function(sub) {
    valid <- unique(sub$time2[sub$occurred == 1L])
    sub[sub$time2 %in% valid, ]
  }
  pe_init <- keep_strata(pe_init)
  pe_cont <- keep_strata(pe_cont)

  n_init <- sum(pe$occurred[pe$SndInitiation == 0])
  n_cont <- sum(pe$occurred[pe$SndInitiation == 1])

  # Fit helper: routes to the right estimator --------------------------------
  .fit_partition <- function(pe_sub, formula, label) {
    n_ev <- sum(pe_sub$occurred)
    if (n_ev < min_events) {
      warning(sprintf("[fit_rem_sep] %s sub-model skipped: only %d focal event(s).",
                      label, n_ev))
      return(NULL)
    }
    if (method == "cox")
      return(survival::coxph(formula, data = pe_sub))

    covar_formula <- .surv_to_covar_formula(formula)
    if (method == "clogit")
      return(.fit_clogit(pe_sub, covar_formula))

    # temporal: deltas from the full edgelist time axis (gaps are sequence-wide)
    times   <- sort(unique(edgelist$time))
    deltas  <- c(times[1L], diff(times))
    n_dyads <- length(all_nodes) * (length(all_nodes) - 1L)
    .fit_temporal(pe_sub, covar_formula, deltas, n_dyads)
  }

  fit_init <- .fit_partition(pe_init, formula_init, "initiation")
  fit_cont <- .fit_partition(pe_cont, formula_cont, "continuation")

  # Log-likelihoods ----------------------------------------------------------
  .ll <- function(fit) {
    if (is.null(fit)) return(c(null = NA_real_, fit = NA_real_))
    if (inherits(fit, "coxph"))
      return(c(null = fit$loglik[1L], fit = fit$loglik[2L]))
    c(null = fit$loglik[1L], fit = fit$loglik[2L])
  }
  ll_i <- .ll(fit_init)
  ll_c <- .ll(fit_cont)

  loglik <- c(
    init = unname(ll_i["fit"]),
    cont = unname(ll_c["fit"]),
    sep  = unname(ll_i["fit"]) + unname(ll_c["fit"]),
    null = unname(ll_i["null"]) + unname(ll_c["null"])
  )

  structure(
    list(
      init          = fit_init,
      cont          = fit_cont,
      loglik        = loglik,
      n_init        = n_init,
      n_cont        = n_cont,
      init_fraction = n_init / (n_init + n_cont),
      method        = method
    ),
    class = "rem_sep"
  )
}

#' @export
print.rem_sep <- function(x, ...) {
  cat(sprintf("Separable REM (%s)\n", x$method))
  cat(sprintf("  Events: %d initiation (%.1f%%)  |  %d continuation\n",
              x$n_init, 100 * x$init_fraction, x$n_cont))
  cat(sprintf("  Log-lik: init = %.1f  cont = %.1f  sep = %.1f\n\n",
              x$loglik["init"], x$loglik["cont"], x$loglik["sep"]))

  if (!is.null(x$init)) {
    cat("-- Initiation sub-model --\n")
    print(x$init, ...)
  } else {
    cat("-- Initiation sub-model: NULL (too few events) --\n")
  }
  cat("\n")
  if (!is.null(x$cont)) {
    cat("-- Continuation sub-model --\n")
    print(x$cont, ...)
  } else {
    cat("-- Continuation sub-model: NULL (too few events) --\n")
  }
  invisible(x)
}
