# Copyright (c) 2024-2025 Ying Wang, Min Li
# SPDX-License-Identifier: GPL-3.0-or-later

#' Degree-of-freedom calculations
#'
#' Mirrors: fastLPR/utility/core/fastLPR_dof.m
#'
#' @keywords internal
fastlpr_dof <- function(regs, weights_list, sum_weights_list) {
  Tx <- regs$Tx
  dh <- length(weights_list)

  df <- numeric(dh)
  diagS_mat <- matrix(0, nrow = Tx, ncol = dh)
  pdof_vec <- numeric(dh)

  for (j in seq_len(dh)) {
    weights <- weights_list[[j]]
    row_sum <- sum_weights_list[[j]]
    diagS <- diag(weights) / row_sum
    df[j] <- sum(diagS)
    diagS_mat[, j] <- diagS
    pdof_vec[j] <- mean(1 - diagS)
  }

  regs$df_m <- df
  regs$dof <- df[which.min(regs$gcv_yhat$gcv_m)]
  regs$diagS <- diagS_mat
  regs$pdof_m <- pdof_vec
  regs
}

#' Estimate degrees of freedom using Hutchinson's randomized trace estimator
#'
#' This is MUCH faster than the naive O(N^2) approach because:
#' 1. No need to compute the full NxN weight matrix
#' 2. Uses existing fast NUFFT-based regression (O(N + M log M))
#' 3. Only 5-10 random samples needed for stable estimate
#'
#' COMPLEXITY:
#'   Naive approach: O(N^2 * dh)  [NxN matrix for each bandwidth]
#'   Hutchinson:     O((N + M log M) * dh * num_samples)  [just M=100-500 typically]
#'   Speedup:        50-100x for large N (N=5000+)
#'
#' Algorithm (from MATLAB fastLPR_dof.m):
#' 1. Generate random vectors: p ~ N(0,I)
#' 2. Compute fitted values: p_hat = H*p (using fast method!)
#' 3. Estimate trace: tr(I-H) ?(1/N) * p'*(p - p_hat) / (p'*p)
#' 4. Repeat for multiple samples and compute statistics
#'
#' VECTORIZATION (2025-11-21):
#' - Process ALL samples simultaneously (no loop!)
#' - Treats num_samples as multi-response dimension
#' - 3-5x speedup vs sequential processing
#'
#' Mirrors: fastLPR/utility/core/fastLPR_dof.m
#'
#' @param regs Regression structure from fastlpr_s
#' @param num_samples Number of random samples (default: 10)
#'
#' @return List with DOF estimates for each bandwidth
#' @keywords internal
fastlpr_dof_hutchinson <- function(regs, num_samples = 10) {

  Tx <- regs$Tx
  dh <- regs$dh
  dy <- regs$dy

  # ==========================================================================
  # Support external random vectors for exact reproducibility
  # If regs$opt$dof_random_vectors is provided (e.g., from MATLAB reference file),
  # use those exact vectors instead of generating new ones.
  # This enables exact matching between R and MATLAB DOF computations.
  # ==========================================================================
  if (!is.null(regs$opt$dof_random_vectors)) {
    p <- regs$opt$dof_random_vectors
    # Validate dimensions
    if (nrow(p) != Tx || ncol(p) != num_samples) {
      stop(sprintf(
        "dof_random_vectors must be (%d x %d), got (%d x %d)",
        Tx, num_samples, nrow(p), ncol(p)
      ))
    }
  } else {
    # Fallback: Generate random probe vectors for DOF estimation
    # NOTE: R's rnorm will NOT match MATLAB's randn because:
    #   - R uses inverse-CDF method, MATLAB uses Ziggurat algorithm
    #   - Even RNGkind("Mersenne-Twister") doesn't help (different seed init)
    # For cross-language reproducibility, use opt$dof_random_vectors from MATLAB.
    # For R-only use, statistical equivalence is sufficient (same N(0,1) distribution).
    if (!is.null(regs$opt$seed)) set.seed(regs$opt$seed)
    p <- matrix(rnorm(Tx * num_samples), nrow = Tx, ncol = num_samples)
  }

  # ===== VECTORIZED SAMPLE PROCESSING =====
  # Process ALL samples at once by treating them as multi-response data
  # This matches MATLAB's approach: mq = fastLPR_reg(regs, p)
  #
  # Key insight: fastlpr_reg interprets multi-column y as multiple responses (dy dimension)
  # Passing p with shape (Tx, num_samples) processes ALL samples in ONE call!
  # Output: mq with shape (Ng, dh, num_samples) - exactly what we need
  #
  # OLD APPROACH (SLOW): Loop over num_samples, call fastlpr_reg separately
  # NEW APPROACH (FAST): Single call, broadcast across samples automatically

  # CRITICAL: Use the ORIGINAL regs$s (design matrix) that was already computed!
  # The design matrix depends on kernel weights at data locations, NOT on response values.
  # Recomputing s for random vectors p was the bug causing 47x underestimation.
  #
  # MATLAB (correct): mq = fastLPR_reg(regs, p)  # Uses precomputed regs.s
  # R (now correct): Just call fastlpr_reg with original regs and new response p
  #
  # Set y_isreal=TRUE for DOF computation (DOF does not depend on response values)
  # The random probe vectors p are ALWAYS real, so the output should be real.
  # But regs$y_isreal might be FALSE if original data y was complex.
  # This prevents unnecessary complex computations and allows Rcpp interpolation.
  y_isreal_orig <- regs$y_isreal
  regs$y_isreal <- TRUE  # Force real output for real probe vectors
  mq <- fastlpr_reg(regs, p)
  regs$y_isreal <- y_isreal_orig  # Restore original setting
  
  # Ensure output is real (discard any numerical noise from FFT)
  if (is.complex(mq)) {
    mq <- Re(mq)
  }


  # MATLAB: mq = reshape(mq, [regs.N, regs.dh*regs.opt.num_dof_sample])
  # Reshape mq to 2D: (Ng, dh*num_samples)
  # This allows interpolation to treat each column as a separate response
  mq_dims <- dim(mq)
  if (is.null(mq_dims)) {
    # Vector: reshape to (Ng, dh*num_samples)
    Ng <- length(mq)
    mq <- matrix(mq, nrow = Ng, ncol = dh * num_samples)
  } else if (length(mq_dims) == 1) {
    # 1D array: reshape to (Ng, dh*num_samples)
    Ng <- mq_dims[1]
    mq <- matrix(mq, nrow = Ng, ncol = dh * num_samples)
  } else if (length(mq_dims) == 2) {
    # Already 2D, but might need reshaping if not (Ng, dh*num_samples)
    # Keep as is if dimensions match, otherwise reshape
    if (ncol(mq) != dh * num_samples) {
      Ng <- nrow(mq)
      mq <- matrix(mq, nrow = Ng, ncol = dh * num_samples)
    }
  } else {
    # 3D or higher: flatten to 2D
    # Use aperm + matrix instead of loop
    # R's fastlpr_reg returns 3D array (Ng, dh, num_samples)
    # MATLAB returns 2D (flattened) with columns ordered as (dh1, dh2, ..., dhN, sample1, sample2, ...)
    # We need to convert (Ng, dh, num_samples) -> (Ng, dh*num_samples)
    Ng <- mq_dims[1]
    dh_actual <- mq_dims[2]
    num_samples_actual <- mq_dims[3]

    # FAST: aperm reorders dimensions, then matrix flattens last two dims
    # aperm(mq, c(1, 2, 3)) keeps same order: (Ng, dh, num_samples)
    # matrix(..., nrow=Ng) flattens (dh, num_samples) -> (dh*num_samples) in column-major order
    # This gives columns: [dh1,s1], [dh2,s1], ..., [dhN,s1], [dh1,s2], [dh2,s2], ...
    mq <- matrix(mq, nrow = Ng, ncol = dh_actual * num_samples_actual)
  }

  # Now mq is 2D: (Ng, dh*num_samples)

  # ===== END VECTORIZED SECTION =====

  # ===== INTERPOLATION: RCPP vs PURE R =====
  # MATLAB uses griddedInterpolant which handles all columns at once
  # R's approxfun requires looping (slow!)
  # Solution: Use Rcpp batch interpolation if available

  use_rcpp <- FALSE
  if (!is.null(regs$opt$use_rcpp)) {
    use_rcpp <- regs$opt$use_rcpp && requireNamespace("Rcpp", quietly = TRUE)
  } else {
    # Auto-detect: Use Rcpp if available and many bandwidths
    # Use get0() instead of exists() because exists() doesn't
    # reliably find functions in the package namespace when called from within
    dof_interp_batch_fn <- tryCatch({
      fn <- get0("dof_interpolate_batch", envir = globalenv(), inherits = FALSE)
      if (is.null(fn)) fn <- get0("dof_interpolate_batch", envir = asNamespace("fastlpr"), inherits = FALSE)
      fn
    }, error = function(e) NULL)
    use_rcpp <- (dh * num_samples > 50) && is.function(dof_interp_batch_fn)
  }

  if (regs$dx == 1) {
    # 1D case
    # MATLAB approach:
    # 1. mq is 2D: (Ng, dh*num_samples)
    # 2. Interpolate each column separately
    # 3. phat_flat is 2D: (Tx, dh*num_samples)
    # 4. Reshape phat_flat to 3D: (Tx, dh, num_samples)
    # 5. Permute to (Tx, num_samples, dh)

    grid_raw <- as.numeric(regs$xlist[[1]])  # Grid points in RAW space
    x_raw <- as.numeric(regs$xraw[, 1])      # Data points in RAW space

    # OPTIMIZATION: Use unified rcpp_interp_batch_nd (N-D, OpenMP parallelized)
    # mq: (Ng, dh * num_samples), output: (Tx, dh * num_samples)
    dims <- as.integer(c(length(grid_raw), ncol(mq)))
    phat_flat <- rcpp_interp_batch_nd(
      list(as.numeric(grid_raw)),
      as.numeric(mq),
      dims,
      matrix(x_raw, ncol = 1)
    )

    # MATLAB: phat = reshape(phat, [regs.Ty, regs.dh, regs.opt.num_dof_sample])
    # Reshape to 3D: (Tx, dh, num_samples)
    phat <- array(phat_flat, dim = c(Tx, dh, num_samples))
    
    # MATLAB: phat = permute(phat, [1,3,2])
    # Permute to (Tx, num_samples, dh)
    phat <- aperm(phat, c(1, 3, 2))  # aperm is R's permute
    
  } else {
    # Multi-dimensional case - more complex interpolation needed
    stop("Hutchinson DOF estimator not yet implemented for dx > 1")
  }

  # Compute pdof using Hutchinson's estimator
  # MATLAB:
  #   phat: (Ty, num_samples, dh) after permute
  #   p: (Ty, num_samples)
  #   pdof = sum(p.*(p-phat)) ./ sum(p.*p)  # Broadcasting and sum along dim 1
  #   pdof: (1, num_samples, dh)
  #   pdof = permute(pdof, [2,3,1]) -> (num_samples, dh, 1)
  
  # R implementation (OPTIMIZED):
  # phat: (Tx, num_samples, dh)
  # p: (Tx, num_samples)
  
  # OPTIMIZATION: Use array recycling instead of explicit loop
  # R automatically recycles p along the 3rd dimension when dimensions match
  p_broadcast <- array(p, dim = c(Tx, num_samples, dh))
  
  # OPTIMIZATION: Use colSums instead of apply(..., sum) - 15x faster
  # For 3D array, we need to collapse first dimension
  # Reshape to 2D: (Tx, num_samples*dh), apply colSums, then reshape back
  diff_array <- p_broadcast * (p_broadcast - phat)
  dim_diff <- dim(diff_array)
  diff_2d <- matrix(diff_array, nrow = dim_diff[1], ncol = dim_diff[2] * dim_diff[3])
  numerator_flat <- colSums(diff_2d)
  numerator <- matrix(numerator_flat, nrow = num_samples, ncol = dh)
  
  # OPTIMIZATION: colSums for denominator
  denominator <- colSums(p * p)  # (num_samples,)
  
  # Broadcast denominator to (num_samples, dh) for division
  denominator_broadcast <- matrix(rep(denominator, dh), nrow = num_samples, ncol = dh)

  # Compute pdof: (num_samples, dh)
  pdof <- numerator / denominator_broadcast

  # Compute df: df = Tx * (1 - pdof)
  df <- Tx * (1 - pdof)


  # Clamp pdof to valid range [0, 1]
  # Hutchinson estimator has variance and can give values outside [0,1]
  # NOTE: pmax/pmin destroy matrix dimensions, so we need to restore them
  pdof_dims <- dim(pdof)
  pdof <- pmax(0, pmin(1, pdof))
  dim(pdof) <- pdof_dims  # Restore matrix dimensions

  # Complex-valued responses need 2x DOF adjustment
  # MATLAB code:
  #   if ~regs.y_isreal
  #       pdof(:,:,~regs.y_isreal)=2*pdof(:,:,~regs.y_isreal)-1;
  #   end
  # This is because complex-valued data has 2 degrees of freedom per observation
  if (!is.null(regs$y_isreal) && !all(regs$y_isreal)) {
    # Apply 2x DOF adjustment for complex responses
    pdof <- 2 * pdof - 1
    # Re-clamp after adjustment
    pdof <- pmax(0, pmin(1, pdof))
    dim(pdof) <- pdof_dims
  }

  # Compute df from CLAMPED pdof
  # When pdof > 1.0, this caused negative df values - compute df AFTER clamping
  df <- Tx * (1 - pdof)

  # Ensure 2D matrices for colMeans/apply (handles dh=1 case)
  pdof <- as.matrix(pdof)
  df <- as.matrix(df)

  # Compute statistics across samples
  pdof_m <- colMeans(pdof)
  pdof_sd <- apply(pdof, 2, sd)
  df_m <- colMeans(df)
  df_sd <- apply(df, 2, sd)

  # ==========================================================================
  # Compute GCV penalty and its standard deviation
  # CRITICAL: Must compute penalty samples FIRST, then take mean/sd
  # The delta method approximation was causing incorrect gcv_sd values
  # ==========================================================================

  # Clamp pdof samples for penalty computation (prevents division by zero)
  # pdof is (num_samples, dh) matrix
  pdof_for_penalty <- pmax(pdof, 0.01)  # Minimum 1% residual DOF

  # Compute penalty samples: penalty = 1 / pdof^2
  penalty_samples <- 1 / (pdof_for_penalty^2)  # (num_samples, dh) matrix

  # Compute mean and std of penalty (matches MATLAB exactly)
  pdof_inv_m <- colMeans(penalty_samples)  # Mean penalty for each bandwidth
  pdof_inv_sd <- apply(penalty_samples, 2, sd)  # Std of penalty for each bandwidth

  return(list(
    pdof_m = pdof_m,
    pdof_sd = pdof_sd,
    df_m = df_m,
    df_sd = df_sd,
    pdof_inv_m = pdof_inv_m,
    pdof_inv_sd = pdof_inv_sd,
    num_samples = num_samples
  ))
}
