#' cox.rvph (Remedy for Violations of the Proportional Hazards Assumption of Cox regression)
#'
#' Stepwise or time-varying remedies for proportional hazards assumption violations of a cox regression model
#'
#' @param data dataset
#' @param time time variable
#' @param event event indicator
#' @param covariate covariate that violates the proportional hazards assumption
#' @param adjust_vars optional adjustment variables
#' @param method "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, inverse, and scaled)
#' is used. Users may also supply custom transformation functions.
#' @param max_K maximum number of segments
#' @param p_threshold threshold for PH test
#' @param verbose logical; whether to print progress messages
#'
#' @return list containing model results
#'
#'
#' @details
#' 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.
#'
#' The \code{step} method performs segmented modeling by searching
#' for optimal cut points in time by maximizing the partial likelihood. 
#' The selected cut points are incorporated into the Cox model
#' through time-dependent interval-specific covariates using the
#' counting-process formulation.
#' If the resulting model satisfies
#' the PH assumption with relatively few cut points, hazard ratios (HRs)
#' may be interpreted within each estimated time interval.
#'
#' However, if many cut points are required or PH violations persist,
#' a smooth time-varying effect may be more appropriate. In such cases,
#' the \code{timev} method fits several candidate time-transformation
#' functions \eqn{g(t)} and selects the model with the smallest AIC.
#' By default, the following candidate functions are evaluated:
#' linear (\eqn{t}),
#' log (\eqn{\log(t+1)}),
#' sqrt (\eqn{\sqrt{t}}),
#' quadratic (\eqn{t^2}),
#' inverse (\eqn{1/(t+1)}),
#' and scaled (\eqn{t/\max(t)}).
#' Users may alternatively provide their own candidate functions
#' through the g_candidates argument.
#'
#' For the selected time-varying model, the hazard ratio at time
#' \eqn{t} is given by:
#'
#' \deqn{
#' HR(t) = \exp(\beta + \gamma g(t))
#' }
#'
#' where \eqn{\beta} is the baseline coefficient 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.
#'
#' Interpretation:
#' For the \code{step} method, hazard ratios are interpreted
#' separately within each estimated time interval.
#'
#' Example output (method: step):
#'
#' \preformatted{
#' $K
#' [1] 2
#'
#' $tau
#' [1] 2.35
#'
#' $p_value
#'
#' [1] 0.7042976
#' 
#' $zph
#'                 chisq df    p
#' covariate_seg1 0.0546  1 0.82
#' covariate_seg2 0.1630  1 0.69
#' age            1.2196  1 0.27
#' GLOBAL         1.4053  3 0.70
#' 
#' $fit
#' Call:
#' survival::coxph(formula = as.formula(f_str_final), data = final_data)
#' 
#'                    coef exp(coef) se(coef)     z        p
#' covariate_seg1 0.90130   2.46281  0.21302 4.231 2.33e-05
#' covariate_seg2 0.06103   1.06293  0.21348 0.286    0.775
#' }
#'
#' Interpretation (method: step):
#'
#' For the \code{step} method, hazard ratios are interpreted
#' separately within each estimated time interval.
#' Interpretation:
#'
#' \verb{covariate_seg1  HR = 2.46}
#'
#' \verb{covariate_seg2  HR = 1.06}
#'
#' If the estimated cut point is \eqn{\tau = 2.35}, this indicates
#' that before time 2.35 the hazard is approximately 2.46 times higher,
#' whereas after time 2.35 the hazard ratio decreases to approximately 1.06.
#'
#' Example output (method: timev):
#'
#' \preformatted{
#' $selected_g
#' [1] "sqrt"
#'
#' $fit
#'
#'                  coef exp(coef)  se(coef)      z        p
#' protime      0.693216  2.000137  0.132824  5.219 1.80e-07
#' tt(protime) -0.013456  0.986634  0.003959 -3.399 0.000677
#' bili         0.118716  1.126050  0.012860  9.231  < 2e-16
#' .
#' .
#'
#' }
#'
#' Interpretation (method: timev):
#'
#' For the \code{timev} method, hazard ratios vary continuously over time.
#'
#' Example model:
#'
#' \deqn{
#' HR(t) = \exp(0.693 - 0.013 \sqrt{t})
#' }
#'
#' This indicates that the hazard ratio decreases gradually 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 logLik quantile
#' @importFrom utils combn
#' @examples
#' if (requireNamespace("KMsurv", quietly = TRUE)) {
#'   data(larynx, package = "KMsurv")
#'
#'   cox.rvph(
#'     data = larynx,
#'     time = "time",
#'     event = "delta",
#'     covariate = "stage",
#'     adjust_vars = "age",
#'     method = "step"
#'   )
#' }
#' @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")
  }
  
  if(any(is.na(data[[time]]))) stop("Missing values in time")
  if(any(is.na(data[[event]]))) stop("Missing values in event")
  
  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]
      
      cuts_i <- tau[tau < ti]
      points <- c(0, cuts_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),
        scaled = function(t) t / max(t)
      )
      
    }
    
    
    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(
      method = "timev",
      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 <- as.numeric(quantile(t[d == 1], probs = seq(0.2, 0.8, by = 0.1)))
    
    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 Tau(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)
}
