#' @keywords internal
.degrees_of_freedom_kr <- function(x, ...) {
  UseMethod(".degrees_of_freedom_kr")
}

#' @keywords internal
.degrees_of_freedom_kr.default <- function(x, ...) {
  # Extract fixed effect parameter names
  parameters <- find_parameters(x, effects = "fixed", flatten = TRUE)

  # Create an identity matrix (L) representing linear contrasts for each parameter
  L <- as.data.frame(diag(rep(1, n_parameters(x, effects = "fixed"))))

  # Calculate the Kenward-Roger adjusted variance-covariance matrix
  krvcov <- .vcov_kenward_ajusted(x)

  # Calculate the denominator degrees of freedom (ddf) for each parameter (each column in L)
  dof <- stats::setNames(
    sapply(L, .kenward_adjusted_ddf, model = x, adjusted_vcov = krvcov),
    parameters
  )

  # Attach standard errors and adjusted vcov as attributes for downstream use
  attr(dof, "vcov") <- krvcov
  attr(dof, "se") <- abs(as.vector(sqrt(diag(as.matrix(krvcov)))))
  dof
}

#' @keywords internal
.degrees_of_freedom_kr.glmmTMB <- function(x, verbose = TRUE, ...) {
  check_if_installed("glmmTMB", minimum_version = "1.1.3")

  # Kenward-Roger relies on Restricted Maximum Likelihood (REML)
  if (!isTRUE(x$modelInfo$REML) && verbose) {
    format_warning(
      "Kenward-Roger degrees of freedom are typically not appropriate for models fit with ML. Use `REML = TRUE`."
    )
  }

  # Use glmmTMB's native KR calculation
  out <- glmmTMB::dof_KR(x)
  krvcov <- as.matrix(attributes(out)$vcov)
  dof <- as.vector(out)

  attr(dof, "vcov") <- krvcov
  attr(dof, "se") <- abs(as.vector(sqrt(diag(krvcov))))
  stats::setNames(dof, names(out))
}


# The following code was taken from the "pbkrtest" package and slightly modified
#' @author Søren Højsgaard, \email{sorenh@@math.aau.dk}
.kenward_adjusted_ddf <- function(model, linear_coef, adjusted_vcov) {
  # Wrapper to pass the unadjusted vcov alongside the adjusted one
  .adjusted_ddf(adjusted_vcov, linear_coef, stats::vcov(model))
}


#' @keywords internal
.adjusted_ddf <- function(adjusted_vcov, linear_coef, unadjusted_vcov = adjusted_vcov) {
  check_if_installed("Matrix")

  # Ensure linear contrast is a column vector
  if (!is.matrix(linear_coef)) {
    linear_coef <- matrix(linear_coef, ncol = 1)
  }

  # Calculate scaling factor 'vlb'
  vlb <- sum(linear_coef * (unadjusted_vcov %*% linear_coef))

  # Matrix theta based on the outer product of the contrast vector
  theta <- Matrix::Matrix(
    as.numeric(outer(linear_coef, linear_coef) / vlb),
    nrow = length(linear_coef)
  )

  # Retrieve derivative matrices computed during the adjusted vcov step
  P <- attr(adjusted_vcov, "P")
  W <- attr(adjusted_vcov, "W")

  A1 <- A2 <- 0
  theta_unadjusted_vcov <- theta %*% unadjusted_vcov
  n_variance_components <- length(P)

  # Taylor expansion approximation steps based on Kenward & Roger (1997)
  # Loops over variance parameters to accumulate adjustment terms A1 and A2
  for (ii in 1:n_variance_components) {
    for (jj in ii:n_variance_components) {
      # Adjust weighting based on whether indices match (diagonal vs off-diagonal)
      e <- ifelse(ii == jj, 1, 2)

      # ui and uj are inner matrix products used to calculate traces
      ui <- as.matrix(theta_unadjusted_vcov %*% P[[ii]] %*% unadjusted_vcov)
      uj <- as.matrix(theta_unadjusted_vcov %*% P[[jj]] %*% unadjusted_vcov)

      # Sum of diag() represents the matrix trace
      A1 <- A1 + e * W[ii, jj] * (sum(diag(ui)) * sum(diag(uj)))
      A2 <- A2 + e * W[ii, jj] * sum(ui * t(uj))
    }
  }

  # Constants and intermediate variables for the final ddf formula
  B <- (A1 + 6 * A2) / 2
  g <- (2 * A1 - 5 * A2) / (3 * A2)
  c1 <- g / (3 + 2 * (1 - g))
  c2 <- (1 - g) / (3 + 2 * (1 - g))
  c3 <- (3 - g) / (3 + 2 * (1 - g))

  # Further adjustment factors
  V0 <- 1 + c1 * B
  V1 <- 1 - c2 * B
  V2 <- 1 - c3 * B
  V0 <- ifelse(abs(V0) < 1e-10, 0, V0)

  # Calculate rho and finally the degrees of freedom (df2)
  rho <- (.divZero(1 - A2, V1))^2 * V0 / V2
  df2 <- 4 + 3 / (rho - 1)
  df2
}


.divZero <- function(x, y, tol = 1e-14) {
  # Safe division to handle edge cases where ratio approaches 1
  if (abs(x) < tol && abs(y) < tol) {
    1
  } else {
    x / y
  }
}


.vcov_kenward_ajusted <- function(model) {
  check_if_installed("lme4")

  # Force REML estimation if the model was fit with ML
  if (!(lme4::getME(model, "is_REML"))) {
    model <- stats::update(model, . ~ ., REML = TRUE)
  }

  # Compute the adjusted covariance matrix
  .vcovAdj16_internal(stats::vcov(model), .get_SigmaG(model), lme4::getME(model, "X"))
}


.get_SigmaG <- function(model) {
  check_if_installed("lme4")
  check_if_installed("Matrix")

  GGamma <- lme4::VarCorr(model)
  re_structure <- .shgetME(model)

  # Combine covariance parameters for all random effects into a single vector (ggamma)
  ggamma <- NULL
  for (ii in 1:(re_structure$n.RT)) {
    Lii <- GGamma[[ii]]
    # Extract lower triangular part including diagonal
    ggamma <- c(ggamma, Lii[lower.tri(Lii, diag = TRUE)])
  }
  # Append the residual variance
  ggamma <- c(ggamma, stats::sigma(model)^2)
  n_variance_components <- length(ggamma)

  # Construct the corresponding G matrices for random effects
  G <- NULL
  Zt <- lme4::getME(model, "Zt")
  for (ss in 1:re_structure$n.RT) {
    ZZ <- .shget_Zt_group(ss, Zt, re_structure$Gp)
    n.lev <- re_structure$n.lev.by.RT2[ss]
    Ig <- Matrix::sparseMatrix(1:n.lev, 1:n.lev, x = 1)

    for (rr in 1:re_structure$n.parm.by.RT[ss]) {
      ii.jj <- .index2UpperTriEntry(rr, re_structure$n.comp.by.RT[ss])
      ii.jj <- unique(ii.jj)

      # Build sparse indicator matrices
      if (length(ii.jj) == 1) {
        EE <- Matrix::sparseMatrix(
          ii.jj,
          ii.jj,
          x = 1,
          dims = rep(re_structure$n.comp.by.RT[ss], 2)
        )
      } else {
        EE <- Matrix::sparseMatrix(
          ii.jj,
          ii.jj[2:1],
          dims = rep(re_structure$n.comp.by.RT[ss], 2)
        )
      }

      # Kronecker product of Identity and EE
      EE <- Ig %x% EE
      G <- c(G, list(t(ZZ) %*% EE %*% ZZ))
    }
  }

  # Append the identity matrix for the residual component
  n.obs <- n_obs(model)
  G <- c(G, list(Matrix::sparseMatrix(1:n.obs, 1:n.obs, x = 1)))

  # Construct Sigma by multiplying the variance parameters with their G matrices
  Sigma <- ggamma[1] * G[[1]]
  for (ii in 2:n_variance_components) {
    Sigma <- Sigma + ggamma[ii] * G[[ii]]
  }

  list(Sigma = Sigma, G = G, n_variance_components = n_variance_components)
}


.index2UpperTriEntry <- function(k, N) {
  # Translates a linear 1D index (k) back into 2D coordinates (i, j)
  # for the upper triangular portion of an N x N symmetric matrix
  aa <- cumsum(N:1)
  aaLow <- c(0, aa[-length(aa)])
  i <- which(aaLow < k & k <= aa)
  j <- k - N * i + N - i * (3 - i) / 2 + i
  c(i, j)
}


.vcovAdj16_internal <- function(Phi, SigmaG, X) {
  check_if_installed("MASS")
  check_if_installed("Matrix")

  # Inverse of the Sigma matrix
  SigmaInv <- chol2inv(chol(Matrix::forceSymmetric(as.matrix(SigmaG$Sigma))))
  n_variance_components <- SigmaG$n_variance_components
  TT <- as.matrix(SigmaInv %*% X)

  # Pre-allocate lists for performance
  HH <- OO <- vector("list", n_variance_components)

  for (ii in 1:n_variance_components) {
    HH[[ii]] <- as.matrix(SigmaG$G[[ii]] %*% SigmaInv)
    OO[[ii]] <- as.matrix(HH[[ii]] %*% X)
  }

  # Find partial derivative matrices PP and QQ
  PP <- vector("list", n_variance_components)
  QQ <- vector("list", (n_variance_components * (n_variance_components + 1)) / 2)

  qq_idx <- 1
  for (rr in 1:n_variance_components) {
    OrTrans <- t(OO[[rr]])
    PP[[rr]] <- Matrix::forceSymmetric(-1 * OrTrans %*% TT)
    for (ss in rr:n_variance_components) {
      QQ[[qq_idx]] <- OrTrans %*% SigmaInv %*% OO[[ss]]
      qq_idx <- qq_idx + 1
    }
  }

  PP <- as.matrix(PP)
  QQ <- as.matrix(QQ)

  # Calculate Ktrace
  Ktrace <- matrix(NA, nrow = n_variance_components, ncol = n_variance_components)
  for (rr in 1:n_variance_components) {
    HrTrans <- t(HH[[rr]])
    for (ss in rr:n_variance_components) {
      Ktrace[rr, ss] <- Ktrace[ss, rr] <- sum(HrTrans * HH[[ss]])
    }
  }

  # Find expected information matrix (IE2)
  IE2 <- matrix(NA, nrow = n_variance_components, ncol = n_variance_components)
  for (ii in 1:n_variance_components) {
    Phi.P.ii <- Phi %*% PP[[ii]]
    for (jj in ii:n_variance_components) {
      www <- .indexSymmat2vec(ii, jj, n_variance_components)
      IE2[ii, jj] <- IE2[jj, ii] <- Ktrace[ii, jj] -
        2 * sum(Phi * QQ[[www]]) +
        sum(Phi.P.ii * (PP[[jj]] %*% Phi))
    }
  }

  # Eigenvalue decomposition to check conditioning
  eigenIE2 <- eigen(IE2, only.values = TRUE)$values
  condi <- min(abs(eigenIE2))

  # Calculate W matrix (inverse of Information matrix), falling back to
  # pseudo-inverse if ill-conditioned
  WW <- if (condi > 1e-10) {
    as.matrix(Matrix::forceSymmetric(2 * solve(IE2)))
  } else {
    as.matrix(Matrix::forceSymmetric(2 * MASS::ginv(IE2)))
  }

  # Calculate U matrix for final adjustment
  UU <- matrix(0, nrow = ncol(X), ncol = ncol(X))
  for (ii in 1:(n_variance_components - 1)) {
    for (jj in (ii + 1):n_variance_components) {
      www <- .indexSymmat2vec(ii, jj, n_variance_components)
      UU <- UU + WW[ii, jj] * (QQ[[www]] - PP[[ii]] %*% Phi %*% PP[[jj]])
    }
  }

  UU <- as.matrix(UU)
  UU <- UU + t(UU) # Make symmetric

  # Add diagonal components
  for (ii in 1:n_variance_components) {
    www <- .indexSymmat2vec(ii, ii, n_variance_components)
    UU <- UU + WW[ii, ii] * (QQ[[www]] - PP[[ii]] %*% Phi %*% PP[[ii]])
  }

  # Final Kenward-Roger adjusted covariance matrix (PhiA)
  GGAMMA <- Phi %*% UU %*% Phi
  PhiA <- Phi + 2 * GGAMMA

  # Save intermediate matrices as attributes for DDF calculations
  attr(PhiA, "P") <- PP
  attr(PhiA, "W") <- WW
  attr(PhiA, "condi") <- condi
  PhiA
}


.indexSymmat2vec <- function(i, j, N) {
  # Translates 2D indices (i, j) of a symmetric N x N matrix back
  # into a 1D vector index mapping to its upper triangular elements.
  k <- if (i <= j) {
    (i - 1) * (N - i / 2) + j
  } else {
    (j - 1) * (N - j / 2) + i
  }
}


.shgetME <- function(model) {
  check_if_installed("lme4")

  # Extract various structural properties of the random effects
  Gp <- lme4::getME(model, "Gp")
  n.RT <- length(Gp) - 1
  n.lev.by.RT <- vapply(lme4::getME(model, "flist"), nlevels, numeric(1))
  n.comp.by.RT <- .get.RT.dim.by.RT(model)
  n.parm.by.RT <- (n.comp.by.RT + 1) * n.comp.by.RT / 2
  n.RE.by.RT <- diff(Gp)

  n.lev.by.RT2 <- n.RE.by.RT / n.comp.by.RT

  list(
    Gp = Gp,
    n.RT = n.RT,
    n.lev.by.RT = n.lev.by.RT,
    n.comp.by.RT = n.comp.by.RT,
    n.parm.by.RT = n.parm.by.RT,
    n.RE.by.RT = n.RE.by.RT,
    n.lev.by.RT2 = n.lev.by.RT2,
    n_rtrms = lme4::getME(model, "n_rtrms")
  )
}


.shget_Zt_group <- function(ii.group, Zt, Gp, ...) {
  # Subsets the Zt matrix (transposed random effects design matrix)
  # for a specific grouping factor
  zIndex.sub <- (Gp[ii.group] + 1):Gp[ii.group + 1]
  as.matrix(Zt[zIndex.sub, ])
}


.get.RT.dim.by.RT <- function(model) {
  check_if_installed("lme4")

  # Output: dimension (no of columns) of covariance matrix for random terms
  if (inherits(model, "mer")) {
    vapply(model@ST, nrow, numeric(1))
  } else {
    lengths(lme4::getME(model, "cnms"))
  }
}
