#' cox.rvph (Remedy for Violations of the Proportional Hazards Assumption in Cox Proportional Hazards Models)
#'
#' Stepwise or time-varying remedies for proportional hazards assumption violations in Cox Proportional Hazards Models
#'
#' @param data a data frame
#' @param time the survival time variable
#' @param event the event indicator variable coded as
#' 0 = right-censored and 1 = event
#' @param covariate the covariate that violates the proportional hazards assumption
#' @param adjust_vars the variables to be included as adjustment covariates
#' @param method the method to be applied ("step" or "timev")
#' @param g_candidates a list of candidate time functions
#' used in the time-varying coefficient method.
#' If NULL, a default set of candidate functions
#' (linear, log, sqrt, quadratic and inverse)
#' is used. Users may also supply custom transformation functions.
#' If supplied, the user-provided functions replace the default candidate set.
#' 
#' @param max_K the maximum number of segments
#' The actual number of searchable segments may be limited by
#' the number of distinct candidate split points.
#' @param p_threshold The significance threshold used for the PH test
#' @param verbose logical; whether to print progress messages
#'
#' @return a list containing the fitted model and related results
#'
#'
#' @details
#' The event indicator specified by \code{event} must be binary,
#' with 0 indicating a right-censored observation and 1 indicating
#' that the event of interest occurred.
#' Other event codings must be recoded before using \code{cox.rvph()}.
#' 
#' Additionally, users should verify that the survival time variable
#' does not contain negative values. For the \code{step} method,
#' observations with a survival time of 0 should be handled using an
#' appropriate preprocessing strategy based on the scientific context.
#' For the \code{timev} method, the specified time-transformation
#' functions should be defined over the observed range of survival times.
#'
#' Users should first assess the proportional hazards (PH) assumption
#' using \code{cox.zph()} before applying \code{cox.rvph()}.
#' Variables showing evidence of non-proportional hazards may then be modeled using stepwise or time-varying remedies.
#'
#' Currently, only a single continuous variable can be specified in
#' \code{covariate}. Categorical variables may still be included
#' in \code{adjust_vars} as adjustment covariates.
#'
#' The \code{step} method performs segmented modeling by searching
#' for optimal split points in time by maximizing the partial likelihood. 
#' The selected split points are incorporated into the Cox model
#' through time-dependent interval-specific covariates using the
#' counting-process formulation.
#' For each number of segments, the PH assumption is assessed using
#' the global test from \code{cox.zph()}.
#' The search stops at the first model for which the global p-value
#' exceeds \code{p_threshold}.
#' If the resulting model satisfies
#' the PH assumption with relatively few split points, hazard ratios (HRs)
#' may be interpreted within each estimated time interval.
#' 
#' The resulting hazard ratio within each time interval \eqn{k}
#' is given by:
#'
#' \deqn{
#' HR(t) = \exp(\beta_k), \quad t \in I_k
#' }
#'
#' However, if many split points are required or PH violations persist,
#' a smooth time-varying effect may be more appropriate. In such cases,
#' the \code{timev} method compares the base Cox model with
#' several time-varying alternatives based on candidate
#' time-transformation functions using AIC, and selects the one
#' with the smallest AIC.
#' By default, the following candidate functions are evaluated:
#' linear (\eqn{t}),
#' logarithmic (log) (\eqn{\log(t+1)}),
#' square-root (sqrt) (\eqn{\sqrt{t}}),
#' quadratic (\eqn{t^2})
#' and inverse (\eqn{1/(t+1)}).
#' Users may alternatively provide their own candidate functions
#' through the \code{g_candidates} argument.
#'
#' The resulting hazard ratio at time \eqn{t} is given by:
#'
#' \deqn{
#' HR(t) = \exp(\beta + \gamma g(t))
#' }
#'
#' where \eqn{\beta} is the coefficient of the covariate and
#' \eqn{\gamma} represents the time-varying interaction effect.
#' Users may evaluate this expression at clinically relevant time points
#' to interpret how the hazard ratio changes over time.
#'
#'
#' Example output (method: step):
#'
#' \preformatted{
#' $K
#' [1] 2
#'
#' $tau
#' [1] 24.8
#'
#' $p_value
#'
#' [1] 0.2529169
#' 
#' $zph
#'                 chisq df    p
#' age_seg1       0.0813  1 0.78
#' age_seg2       3.8528  1 0.05
#' sex            0.3214  1 0.57
#' GLOBAL         4.0804  3 0.25
#' 
#' $fit
#' Call:
#' survival::coxph(formula = as.formula(f_str_final), data = final_data)
#' 
#'               coef exp(coef) se(coef)     z        p
#' age_seg1  0.27052   1.31065  0.07944  3.405 0.000661
#' age_seg2  0.05550   1.05707  0.09677  0.574 0.566284
#' .
#' .
#' }
#'
#' Interpretation (method: step):
#'
#' For the \code{step} method, hazard ratios are interpreted
#' separately within each estimated time interval.
#'
#' \deqn{
#' HR(t) =
#' \begin{cases}
#' \exp(0.27052) = 1.31, & t \le 24.8, \\
#' \exp(0.05550) = 1.06, & t > 24.8.
#' \end{cases}
#' }
#'
#' If the estimated split point is \eqn{\tau = 24.8}, this indicates
#' that the hazard ratio associated with a one-unit increase in age
#' is approximately 1.31 before time 24.8 and decreases to
#' approximately 1.06 after time 24.8.
#'
#' Example output (method: timev):
#'
#' \preformatted{
#' $selected_g
#' [1] "quadratic"
#' 
#' $AIC
#'     base    linear       log      sqrt quadratic   inverse    
#' 1151.107  1139.675  1147.434  1142.833  1138.358  1152.789  
#' 
#' $fit
#' Call:
#' survival::coxph(formula = as.formula(f_str), data = data, tt = tt_fun)
#' 
#'                coef exp(coef)  se(coef)     z        p
#' bili      8.036e-02 1.084e+00 1.807e-02 4.448 8.66e-06
#' tt(bili)  3.563e-08 1.000e+00 8.323e-09 4.281 1.86e-05
#' ascites1  1.153e+00 3.168e+00 2.746e-01 4.199 2.68e-05
#' .
#' .
#'
#' }
#'
#' Interpretation (method: timev):
#'
#' For the \code{timev} method, hazard ratios vary continuously over time.
#'
#' \deqn{
#' HR(t) = \exp(0.08036 + 3.563 \times 10^{-8} t^2)
#' }
#'
#' In this example, the quadratic time transformation indicates that
#' the effect of \code{bili} changes continuously over time.
#' Users may substitute clinically meaningful values of \eqn{t}
#' to estimate hazard ratios at specific time points.
#'
#' @importFrom survival coxph cox.zph Surv
#' @importFrom stats AIC as.formula complete.cases logLik quantile
#' @importFrom utils combn
#' @examples
#' if (requireNamespace("KMsurv", quietly = TRUE)) {
#'   data(psych, package = "KMsurv")
#'   
#'   psych$sex <- factor(psych$sex)
#'
#'   cox.rvph(
#'     data = psych,
#'     time = "time",
#'     event = "death",
#'     covariate = "age",
#'     adjust_vars = "sex",
#'     method = "step"
#'   )
#' }
#' 
#' if (requireNamespace("survival", quietly = TRUE)) {
#'   data(pbc, package = "survival")
#'   
#'   pbc$status2 <- ifelse(pbc$status == 2, 1, 0)
#'   pbc$ascites <- factor(pbc$ascites)
#'   
#'   cox.rvph(
#'     data = pbc,
#'     time = "time",
#'     event = "status2",
#'     covariate = "bili",
#'     adjust_vars = c("ascites", "edema", "protime"),
#'     method = "timev"
#'   )
#' }
#' 
#' 
#' 
#' @export
cox.rvph <- function(data, time, event,
                     covariate, adjust_vars = NULL,
                     method = c("step", "timev"),
                     g_candidates = NULL,
                     max_K = 4, p_threshold = 0.05, verbose = TRUE){
  
  method <- match.arg(method)
  
  if(!all(c(time, event, covariate) %in% names(data))) {
    stop("time/event/covariate not found in data")
  }
  
  ## Check adjustment variables
  if(!is.null(adjust_vars) &&
     !all(adjust_vars %in% names(data))) {
    stop("Some adjust_vars were not found in data")
  }
  
  ## Check event coding
  if(!is.numeric(data[[event]]) &&
     !is.integer(data[[event]])) {
    stop("event must be numeric and coded as 0 = right-censored and 1 = event")
  }
  
  event_values <- unique(data[[event]][!is.na(data[[event]])])
  
  if(!all(event_values %in% c(0, 1))) {
    stop("event must be coded as 0 = right-censored and 1 = event")
  }
  
  ## Variables used in the analysis
  analysis_vars <- unique(c(
    time,
    event,
    covariate,
    adjust_vars
  ))
  
  ## Complete-case analysis
  complete_idx <- complete.cases(data[, analysis_vars, drop = FALSE])
  
  n_removed <- sum(!complete_idx)
  
  if(n_removed > 0 && verbose) {
    message(
      n_removed,
      " observations with missing values were removed"
    )
  }
  
  data <- data[complete_idx, , drop = FALSE]
  
  if(nrow(data) == 0) {
    stop("No complete observations remain after removing missing values")
  }
  
  if(sum(data[[event]] == 1) == 0) {
    stop("No events remain after removing missing values")
  }
  
  make_split_data <- function(data, time, event, covariate, tau) {
    
    tau <- sort(tau)
    K <- length(tau) + 1
    
    split_list <- lapply(seq_len(nrow(data)), function(i) {
      
      ti <- data[[time]][i]
      di <- data[[event]][i]
      xi <- data[[covariate]][i]
      
      splits_i <- tau[tau < ti]
      points <- c(0, splits_i, ti)
      
      temp <- data[rep(i, length(points) - 1), , drop = FALSE]
      
      temp$tstart <- points[-length(points)]
      temp$tstop <- points[-1]
      temp$event_split <- 0
      temp$event_split[nrow(temp)] <- di
      
      seg_id <- findInterval(temp$tstart, vec = tau) + 1
      
      for(k in seq_len(K)) {
        temp[[paste0(covariate, "_seg", k)]] <- ifelse(seg_id == k, xi, 0)
      }
      
      temp
    })
    
    do.call(rbind, split_list)
  }
  
  if(method == "timev") {
    
    if(is.null(g_candidates)) {
      
      g_candidates <- list(
        linear = function(t) t,
        log = function(t) log(t + 1),
        sqrt = function(t) sqrt(t),
        quadratic = function(t) t^2,
        inverse = function(t) 1/(t + 1)
      )
      
    }
    
    
    base_formula <- paste0(
      "survival::Surv(", time, ",", event, ") ~ ",
      covariate,
      if(!is.null(adjust_vars))
        paste(" +", paste(adjust_vars, collapse = " + ")) else ""
    )
    
    base_fit <- survival::coxph(as.formula(base_formula), data = data)
    
    base_aic <- AIC(base_fit)
    
    aic_values <- c(base = base_aic)
    fit_list <- list(base = base_fit)
    
    for(i in seq_along(g_candidates)){
      
      g_fun <- g_candidates[[i]]
      
      tt_fun <- local({
        gf <- g_fun
        
        function(x, t, ...) {
          x * gf(t)
        }
      })
      
      f_str <- paste0(
        "survival::Surv(", time, ",", event, ") ~ ",
        covariate, " + tt(", covariate, ")",
        if(!is.null(adjust_vars))
          paste(" +", paste(adjust_vars, collapse = " + ")) else ""
      )
      
      fit.current <- survival::coxph(
        as.formula(f_str),
        data = data,
        tt = tt_fun
      )
      
      aic_values[i + 1] <- AIC(fit.current)
      names(aic_values)[i + 1] <- names(g_candidates)[i]
      
      fit_list[[i + 1]] <- fit.current
    }
    
    best_idx <- which.min(aic_values)
    
    best_fit <- fit_list[[best_idx]]
    
    best_g <- names(aic_values)[best_idx]
    
    if(verbose){
      message("Time-varying method applied")
      message(paste("Selected g(t):", best_g))
    }
    
    return(list(
      selected_g = best_g,
      AIC = aic_values,
      fit = best_fit
    ))
    
  }
  
  final_result <- NULL
  
  for (K in 2:max_K) {
    if(verbose){
      message("-------------------------------------------")
      message(paste("\nCurrent K =", K, "(number of segments) in progress..."))
    }
    
    t <- data[[time]]
    d <- data[[event]]
    x <- data[[covariate]]
    
    tau.cand <- unique(as.numeric(
      quantile(t[d == 1], probs = seq(0.2, 0.8, by = 0.1))
    ))
    
    if (K - 1 > length(tau.cand)) break
    
    tau.grid <- combn(tau.cand, K - 1, simplify = FALSE)
    
    get_logPL <- function(tau) {
      
      temp_data <- make_split_data(
        data = data,
        time = time,
        event = event,
        covariate = covariate,
        tau = tau
      )
      
      seg_names <- paste0(covariate, "_seg", 1:K)
      
      seg_id_event <- findInterval(temp_data$tstart, vec = tau) + 1
      
      events_per_seg <- tapply(temp_data$event_split, seg_id_event, sum)
      
      if(any(events_per_seg < 5)) return(-Inf)
      
      f_str <- paste0(
        "survival::Surv(tstart, tstop, event_split) ~ ",
        paste(seg_names, collapse = " + "),
        if(!is.null(adjust_vars))
          paste(" +", paste(adjust_vars, collapse = " + "))
        else ""
      )
      
      fit <- try(
        survival::coxph(as.formula(f_str), data = temp_data),
        silent = TRUE
      )
      
      if(inherits(fit, "try-error")) return(-Inf)
      
      return(as.numeric(logLik(fit)))
    }
    
    logPL_list <- sapply(tau.grid, get_logPL)
    
    if(all(is.infinite(logPL_list))) next
    
    best_idx <- which.max(logPL_list)
    tau.hat <- tau.grid[[best_idx]]
    
    final_data <- make_split_data(
      data = data,
      time = time,
      event = event,
      covariate = covariate,
      tau = tau.hat
    )
    
    seg_names <- paste0(covariate, "_seg", 1:K)
    
    f_str_final <- paste0(
      "survival::Surv(tstart, tstop, event_split) ~ ",
      paste(seg_names, collapse = " + "),
      if(!is.null(adjust_vars))
        paste(" +", paste(adjust_vars, collapse = " + "))
      else ""
    )
    
    fit.current <- survival::coxph(as.formula(f_str_final), data = final_data)
    
    zph_test <- survival::cox.zph(fit.current)
    p_val <- zph_test$table["GLOBAL", "p"]
    
    if(verbose) {
      message(paste("Optimal split point(s):", paste(tau.hat, collapse = ", ")))
      message(paste("Global p-value:", round(p_val, 4)))
    }
    
    if (p_val > p_threshold) {
      if(verbose) {
        message(paste("PH assumption satisfied (p >", p_threshold, ")"))
      }
      final_result <- list(K = K, tau = tau.hat, p_value = p_val,zph = zph_test, fit = fit.current)
      return(final_result)
    }
    
    final_result <- list(K = K, tau = tau.hat, p_value = p_val, zph = zph_test, fit = fit.current)
  }
  
  if(verbose) {
    message("PH assumption NOT satisfied up to max_K")
  }
  return(final_result)
}