#' Wiener Process Step-Stress Accelerated Degradation Model Estimation
#'
#' Fits a Wiener process step-stress accelerated degradation model using Maximum Likelihood Estimation (MLE)
#' or Bayesian Markov Chain Monte Carlo (MCMC) as described in Pan and Balakrishnan (2010).
#'
#' @param data A data frame containing inspection times in the first column and unit degradation measurements in subsequent columns,
#'   or a long-format data frame with columns \code{unit}, \code{time}, \code{degradation}.
#' @param stress_levels Vector of stress levels S = (S1, S2, ..., SK).
#' @param thresholds Vector of degradation thresholds (omega_1, ..., omega_K-1) at which stress levels are elevated.
#' @param model Acceleration model: \code{"arrhenius"} or \code{"power"}.
#' @param method Estimation method: \code{"mle"} or \code{"mcmc"}.
#' @param init Initial parameter values \code{c(a, b, sigma)}.
#' @param n_iter Number of MCMC iterations (default: 10000).
#' @param burnin Number of burn-in iterations (default: 2000).
#' @param thin Thinning interval for MCMC (default: 1).
#'
#' @return An object of class \code{"ssadfit"} containing fit results, parameter estimates, log-likelihood, MCMC samples (if applicable), and metadata.
#'
#' @references
#' Pan, Z., & Balakrishnan, N. (2010). Multiple-Steps Step-Stress Accelerated Degradation Modeling Based on Wiener and Gamma Processes.
#' \emph{Communications in Statistics - Simulation and Computation}, 39(7), 1384-1402. \doi{10.1080/03610918.2010.496060}
#'
#' @export
#' @examples
#' data(wiener_arrhenius)
#' fit_w <- ssad_wiener(
#'   data = wiener_arrhenius,
#'   stress_levels = c(45, 65, 85),
#'   thresholds = c(90, 160),
#'   model = "arrhenius",
#'   method = "mle"
#' )
#' print(fit_w)
ssad_wiener <- function(data, stress_levels = c(45, 65, 85), thresholds = c(90, 160),
                        model = c("arrhenius", "power"), method = c("mle", "mcmc"),
                        init = NULL, n_iter = 5000, burnin = 1000, thin = 1) {
  model <- match.arg(model)
  method <- match.arg(method)
  
  # Format data into unit list with times and degradation measurements
  parsed_data <- parse_ssad_data(data)
  units_data <- parsed_data$units
  
  K <- length(stress_levels)
  if (length(thresholds) != K - 1) {
    stop("Length of thresholds must equal length(stress_levels) - 1.")
  }
  
  # Estimate transition times tau_{i, k} for each unit i
  tau_mat <- estimate_ssad_taus(units_data, thresholds)
  
  eval_acc <- function(S, m_type, par) {
    if (is.function(m_type)) return(m_type(S, par))
    a <- par[1]; b <- par[2]
    if (m_type == "arrhenius") return(exp(a + b / (273 + S)))
    if (m_type == "power") return(a * (S^b))
    stop("Unsupported acceleration model.")
  }
  
  # Complete log-likelihood function
  loglik_fn <- function(par) {
    a <- par[1]; b <- par[2]; sigma <- par[3]
    if (is.na(sigma) || sigma <= 1e-6) return(-1e12)
    
    beta_k <- sapply(stress_levels, function(s) eval_acc(s, model, c(a, b)))
    if (any(is.na(beta_k)) || any(beta_k <= 1e-8)) return(-1e12)
    
    logL <- 0
    N <- length(units_data)
    
    # 1. Stress elevation time likelihood g(tau_i)
    d_omega <- c(thresholds[1], diff(thresholds))
    for (i in seq_len(N)) {
      tau_i <- tau_mat[i, ]
      d_tau <- c(tau_i[1], diff(tau_i))
      for (k in seq_len(K - 1)) {
        dtk <- max(d_tau[k], 1e-6)
        dok <- d_omega[k]
        bk <- beta_k[k]
        
        term1 <- 0.5 * log((dok^2) / (2 * pi * (dtk^3) * (sigma^2)))
        term2 <- -((bk * dtk - dok)^2) / (2 * (sigma^2) * dtk)
        logL <- logL + term1 + term2
      }
    }
    
    # 2. Degradation increment likelihood f(dx_{i,j})
    for (i in seq_len(N)) {
      u_df <- units_data[[i]]
      t_vec <- u_df$time
      x_vec <- u_df$degradation
      tau_i <- c(0, tau_mat[i, ], max(t_vec) * 10)
      
      M <- length(t_vec)
      for (j in seq_len(M - 1)) {
        t1 <- t_vec[j]; t2 <- t_vec[j + 1]
        dx <- x_vec[j + 1] - x_vec[j]
        dt <- t2 - t1
        
        # Calculate expected increment dm based on stress intervals
        dm <- 0
        for (k in seq_len(K)) {
          t_start_k <- tau_i[k]
          t_end_k <- tau_i[k + 1]
          overlap <- max(0, min(t2, t_end_k) - max(t1, t_start_k))
          if (overlap > 0) {
            dm <- dm + beta_k[k] * overlap
          }
        }
        
        term1 <- -0.5 * log(2 * pi * dt * (sigma^2))
        term2 <- -((dx - dm)^2) / (2 * (sigma^2) * dt)
        logL <- logL + term1 + term2
      }
    }
    
    if (is.na(logL) || is.nan(logL)) return(-1e12)
    return(logL)
  }
  
  if (is.null(init)) {
    if (model == "arrhenius") {
      init <- c(a = 5.0, b = -2500, sigma = 0.1)
    } else {
      init <- c(a = 1e-3, b = 1.2, sigma = 0.1)
    }
  }
  
  if (method == "mle") {
    neg_loglik <- function(p) -loglik_fn(p)
    opt <- stats::optim(init, neg_loglik, method = "Nelder-Mead", control = list(maxit = 2000))
    opt <- stats::optim(opt$par, neg_loglik, method = "BFGS", hessian = TRUE, control = list(maxit = 2000))
    
    par_mle <- opt$par
    names(par_mle) <- c("a", "b", "sigma")
    
    se_mle <- rep(NA, 3)
    if (!is.null(opt$hessian) && !any(is.na(opt$hessian))) {
      cov_mat <- tryCatch(solve(opt$hessian), error = function(e) NULL)
      if (!is.null(cov_mat)) {
        se_mle <- sqrt(pmax(diag(cov_mat), 0))
      }
    }
    names(se_mle) <- c("a", "b", "sigma")
    
    res <- list(
      par = par_mle,
      se = se_mle,
      loglik = -opt$value,
      process = "wiener",
      model = model,
      method = "mle",
      convergence = opt$convergence,
      stress_levels = stress_levels,
      thresholds = thresholds,
      tau_mat = tau_mat,
      units_data = units_data
    )
  } else { # MCMC
    # Metropolis-Hastings MCMC Sampler
    chain <- matrix(0, nrow = n_iter, ncol = 3)
    colnames(chain) <- c("a", "b", "sigma")
    curr_par <- init
    curr_loglik <- loglik_fn(curr_par)
    
    scale_vec <- abs(init) * 0.05
    scale_vec[scale_vec == 0] <- 0.01
    
    acc_count <- 0
    for (iter in seq_len(n_iter)) {
      prop_par <- curr_par + stats::rnorm(3, mean = 0, sd = scale_vec)
      if (prop_par[3] > 1e-6) {
        prop_loglik <- loglik_fn(prop_par)
        
        # Log prior (normal for a, b, 1/sigma for sigma)
        prior_curr <- stats::dnorm(curr_par[1], 0, 100, log = TRUE) +
                      stats::dnorm(curr_par[2], 0, 10000, log = TRUE) - log(curr_par[3])
        prior_prop <- stats::dnorm(prop_par[1], 0, 100, log = TRUE) +
                      stats::dnorm(prop_par[2], 0, 10000, log = TRUE) - log(prop_par[3])
        
        log_alpha <- (prop_loglik + prior_prop) - (curr_loglik + prior_curr)
        if (!is.na(log_alpha) && log(stats::runif(1)) < log_alpha) {
          curr_par <- prop_par
          curr_loglik <- prop_loglik
          acc_count <- acc_count + 1
        }
      }
      chain[iter, ] <- curr_par
    }
    
    post_chain <- chain[(burnin + 1):n_iter, , drop = FALSE]
    if (thin > 1) {
      post_chain <- post_chain[seq(1, nrow(post_chain), by = thin), , drop = FALSE]
    }
    
    par_mean <- colMeans(post_chain)
    par_sd <- apply(post_chain, 2, stats::sd)
    par_quant <- t(apply(post_chain, 2, stats::quantile, probs = c(0.025, 0.5, 0.975)))
    
    res <- list(
      par = par_mean,
      sd = par_sd,
      quantiles = par_quant,
      chain = post_chain,
      acceptance_rate = acc_count / n_iter,
      loglik = loglik_fn(par_mean),
      process = "wiener",
      model = model,
      method = "mcmc",
      stress_levels = stress_levels,
      thresholds = thresholds,
      tau_mat = tau_mat,
      units_data = units_data
    )
  }
  
  class(res) <- "ssadfit"
  return(res)
}

# Helper functions for data parsing and tau estimation
parse_ssad_data <- function(data) {
  if (is.data.frame(data)) {
    if (all(c("unit", "time", "degradation") %in% colnames(data))) {
      u_names <- unique(data$unit)
      units_list <- lapply(u_names, function(u) {
        sub_d <- data[data$unit == u, ]
        sub_d <- sub_d[order(sub_d$time), ]
        data.frame(time = sub_d$time, degradation = sub_d$degradation)
      })
      names(units_list) <- paste0("unit", u_names)
      return(list(units = units_list))
    } else {
      times <- data[[1]]
      u_cols <- colnames(data)[-1]
      units_list <- lapply(u_cols, function(col) {
        data.frame(time = times, degradation = data[[col]])
      })
      names(units_list) <- u_cols
      return(list(units = units_list))
    }
  } else {
    stop("data must be a data frame.")
  }
}

estimate_ssad_taus <- function(units_data, thresholds) {
  N <- length(units_data)
  K_sub <- length(thresholds)
  tau_mat <- matrix(0, nrow = N, ncol = K_sub)
  colnames(tau_mat) <- paste0("tau", seq_len(K_sub))
  
  for (i in seq_len(N)) {
    u_df <- units_data[[i]]
    t_vec <- u_df$time
    x_vec <- u_df$degradation
    
    for (k in seq_len(K_sub)) {
      om <- thresholds[k]
      idx <- which(x_vec >= om)
      if (length(idx) == 0) {
        # If not crossed, extrapolate based on last slope
        slope <- (max(x_vec) - min(x_vec)) / (max(t_vec) - min(t_vec))
        t_est <- max(t_vec) + (om - max(x_vec)) / max(slope, 1e-4)
      } else if (idx[1] == 1) {
        t_est <- t_vec[1]
      } else {
        # Linear interpolation
        j <- idx[1]
        t1 <- t_vec[j - 1]; t2 <- t_vec[j]
        x1 <- x_vec[j - 1]; x2 <- x_vec[j]
        t_est <- t1 + (om - x1) * (t2 - t1) / max(x2 - x1, 1e-6)
      }
      tau_mat[i, k] <- t_est
    }
  }
  return(tau_mat)
}
