#' Mahalanobis-based cluster-wise regression (internal)
#'
#' Uses an adaptive, regularised covariance matrix per cluster. Also serves
#' as the Mahalanobis clustering solver when `X = NULL` (intercept-only).
#'
#' @noRd
PDVRegressionWithCovar <- function(X = NULL, Y, K, a, b, alpha, beta,
                                   threshold = 0.01, max.iter = 100) {

  if (is.null(X)) {
    X <- matrix(1, nrow(Y), 1)
  } else {
    X <- as.matrix(X)
    X <- cbind(rep(1, nrow(X)), X)
  }
  Y <- as.matrix(Y)

  const <- b / (a - 1)

  dx <- ncol(X)
  dy <- ncol(Y)
  n  <- nrow(X)

  B    <- array(stats::runif(dx * dy * K), dim = c(dx, dy, K))
  newB <- B
  if (dx == 1) {
    B    <- array(t(stats::kmeans(Y, K)$centers), dim = c(dx, dy, K))
    newB <- B
  }

  omega <- array(dim = c(dy, dy, K))
  rho   <- rep(1 / K, K)
  covs  <- array(dim = c(dy, dy, K))

  for (i in 1:K) {
    omega[, , i] <- diag(dy)
  }
  for (i in 1:K) {
    covs[, , i] <- rho[i] * omega[, , i]
  }

  d    <- data.frame(matrix(ncol = K, nrow = n))
  newD <- data.frame(matrix(ncol = K, nrow = n))
  p    <- data.frame(matrix(0, ncol = K, nrow = n))

  for (x in 1:K) {
    names(d)[x] <- paste("Cluster", x)
    names(p)[x] <- paste("Cluster", x)
  }

  JDF     <- c()
  newNorm <- threshold + 1
  iter    <- 0

  while (newNorm >= threshold && iter < max.iter) {
    iter <- iter + 1

    for (j in 1:K) {
      inv <- solve(covs[, , j])
      for (i in 1:n) {
        resids  <- t(matrix(unlist(Y[i, ]))) - t(matrix(unlist(X[i, ]))) %*% B[, , j]
        d[i, j] <- (resids %*% inv %*% t(resids))^(1 / 2)
      }
    }

    JDF_value <- 0
    for (x in 1:n) {
      JDF_numerator   <- prod(d[x, ])^const
      JDF_denominator <- 0
      for (y in 1:K) {
        JDF_denominator <- JDF_denominator + prod(d[x, -y])^const
      }
      JDF_value <- JDF_value + (JDF_numerator / JDF_denominator)^(a - 1)
    }
    JDF <- c(JDF, JDF_value)

    for (x in 1:n) {
      sum1 <- 0
      for (y in 1:K) {
        sum1 <- sum1 + prod(d[x, -y])^const
      }
      for (z in 1:K) {
        product1  <- prod(d[x, -z])^const
        p[x, z]   <- product1 / sum1
      }
    }
    p[sapply(p, is.nan)] <- 1

    for (j in 1:K) {
      ans1 <- matrix(0, ncol = dy, nrow = dx)
      ans2 <- matrix(0, ncol = dx, nrow = dx)
      for (i in 1:n) {
        ans1 <- ans1 + p[i, j]^a * d[i, j]^(b - 2) *
          (matrix(unlist(X[i, ])) %*% t(matrix(unlist(Y[i, ]))))
        ans2 <- ans2 + p[i, j]^a * d[i, j]^(b - 2) *
          (matrix(unlist(X[i, ])) %*% t(matrix(unlist(X[i, ]))))
      }
      newB[, , j] <- solve(ans2) %*% ans1
    }

    for (j in 1:K) {
      invO <- solve(omega[, , j])
      for (i in 1:n) {
        resids     <- t(matrix(unlist(Y[i, ]))) - t(matrix(unlist(X[i, ]))) %*% B[, , j]
        newD[i, j] <- (resids %*% invO %*% t(resids))^(1 / 2)
      }
    }

    S <- array(dim = c(dy, dy, K))
    for (j in 1:K) {
      ans <- matrix(0, ncol = dy, nrow = dy)
      for (i in 1:n) {
        ans1 <- t(matrix(unlist(Y[i, ]))) - t(matrix(unlist(X[i, ]))) %*% B[, , j]
        ans  <- ans + p[i, j]^a * newD[i, j]^(b - 2) * (t(ans1) %*% ans1)
      }
      S[, , j] <- ans
    }

    for (j in 1:K) {
      omega[, , j] <- S[, , j] / det(S[, , j])^(1 / dy)
    }

    sum <- 0
    for (j in 1:K) {
      sum1 <- 0
      for (i in 1:n) {
        sum1 <- sum1 + newD[i, j]^b %*% p[i, j]^a
      }
      sum <- sum + (sum1)^(2 / (b + 2))
    }
    for (j in 1:K) {
      sum2 <- 0
      for (i in 1:n) {
        sum2 <- sum2 + newD[i, j]^b %*% p[i, j]^a
      }
      rho[j] <- (sum2)^(2 / (b + 2)) / sum
    }

    for (j in 1:K) {
      covs[, , j] <- rho[j] * omega[, , j]
    }
    for (j in 1:K) {
      covs[, , j] <- (1 - alpha) * covs[, , j] +
        alpha * det(stats::cov(cbind(X, Y)))^(1 / dy) * diag(dy)
    }

    for (j in 1:K) {
      eigen_decomp <- eigen(omega[, , j])
      eigenvalues  <- eigen_decomp$values
      eigenvectors <- eigen_decomp$vectors
      max_eigen    <- max(eigenvalues)

      for (s in 1:dy) {
        if (max_eigen > beta * eigenvalues[s]) {
          eigenvalues[s] <- max_eigen / beta
        }
      }
      covs[, , j] <- eigenvectors %*% diag(eigenvalues) %*% solve(eigenvectors)
    }

    newNorm <- sqrt(sum(B - newB)^2)
    B <- newB
  }

  return(list(B = B, d = d, p = p, JDF = JDF, rho = rho, cov = covs,
              l = max.col(p)))
}
