#' Density approximation for the DL test statistic
#'
#' Internal helper. This file will contain the almost-exact mixture density for
#' the DerSimonian-Laird test statistic.
#'
#' @keywords internal
aedl_density_tdl <- function(t_grid, K, sig2i, tau2, theta = 0, nsim = 10000,
                             nsim_chi = 100000, n_tau = 1000) {
  if (K < 2) {
    stop("K must be at least 2.", call. = FALSE)
  }

  sim <- sapply(seq_len(K - 1), function(k) {
    stats::rnorm(nsim, mean = theta, sd = sqrt(sig2i[k] + tau2))
  })

  if (K - 1 == 1) {
    sim <- matrix(sim, ncol = 1)
  }

  tau2dl_density <- aedl_density_tau2dl(
    K = K,
    sig2i = sig2i,
    tau2 = tau2,
    nsim_chi = nsim_chi,
    n_tau = n_tau
  )

  tau_grid <- tau2dl_density$x

  gs <- aedl_conditional_moments(
    zi = sim,
    x = tau_grid,
    K = K,
    sig2i = sig2i,
    tau2 = tau2
  )

  gs_x <- rowSums(gs$fztau2)
  gs_x_safe <- gs_x
  gs_x_safe[gs_x_safe == 0] <- NA_real_

  gs_meanx <- gs$fztau2 %*% sim / gs_x_safe
  gs1_meanx <- rowSums(gs$g1) / gs_x_safe

  cov_monte <- lapply(seq_along(tau_grid), function(ix) {
    C <- matrix(0, nrow = K, ncol = K)
    denom_ix <- sum(gs$fztau2[ix, ])

    if (denom_ix == 0 || is.na(denom_ix)) {
      return(C)
    }

    for (ik in seq_len(K)) {
      if (ik != K) {
        vi_center <- sim[, ik] - gs_meanx[ix, ik]
      } else {
        vi_center <- gs$g1[ix, ] - gs1_meanx[ix]
      }

      for (jk in seq_len(K)) {
        if (jk == K && ik == K) {
          C[ik, jk] <- sum(gs$g2[ix, ]) / denom_ix
        } else if (ik == K) {
          vj_center <- sim[, jk] - gs_meanx[ix, jk]
          C[ik, jk] <- sum(vj_center * gs$g1[ix, ]) / denom_ix
        } else if (jk == K) {
          C[ik, jk] <- sum(vi_center * gs$g1[ix, ]) / denom_ix
        } else {
          vj_center <- sim[, jk] - gs_meanx[ix, jk]
          C[ik, jk] <- sum(vi_center * vj_center * gs$fztau2[ix, ]) /
            denom_ix
        }
      }
    }

    C
  })

  v2 <- aedl_v2tdls(
    cov_list = cov_monte,
    x = tau_grid,
    sig2i = sig2i,
    K = K
  )

  bad <- is.nan(v2) | is.na(v2) | v2 <= 0

  if (any(bad)) {
    good <- which(!bad)
    if (length(good) >= 2) {
      v2[bad] <- stats::approx(
        x = good,
        y = v2[good],
        xout = which(bad),
        rule = 2
      )$y
    } else if (length(good) == 1) {
      v2[bad] <- v2[good]
    } else {
      stop("All conditional variance estimates are invalid.", call. = FALSE)
    }
  }

  density <- numeric(length(t_grid))
  dx <- tau_grid[2] - tau_grid[1]

  for (ix in seq_along(tau_grid)) {
    mean_ix <- theta * sqrt(sum(1 / (sig2i + pmax(0, tau_grid[ix]))))

    density <- density +
      stats::dnorm(t_grid, mean = mean_ix, sd = sqrt(v2[ix])) *
      tau2dl_density$y[ix] * dx
  }

  list(
    x = t_grid,
    y = density,
    tau_grid = tau_grid,
    tau_density = tau2dl_density$y,
    v2 = v2
  )
}

#' Conditional Monte Carlo quantities for the DL statistic
#'
#' @keywords internal
aedl_conditional_moments <- function(zi, x, K, sig2i, tau2) {
  s1 <- sum(1 / sig2i)
  s2 <- sum(1 / sig2i^2)
  denom <- s1 - s2 / s1

  len_zi <- nrow(zi)
  len_x <- length(x)

  xmat <- matrix(x, nrow = len_zi, ncol = len_x, byrow = TRUE)

  theta_w <- matrix(
    rowSums(sweep(zi, 2, sig2i[-K], "/")),
    nrow = len_zi,
    ncol = len_x
  )

  theta_w2 <- matrix(
    rowSums(sweep(zi^2, 2, sig2i[-K], "/")),
    nrow = len_zi,
    ncol = len_x
  )

  aK <- 1 - 1 / (s1 * sig2i[K])
  bK <- aK / sig2i[K]

  cval <-
    theta_w^2 / (s1 * sig2i[K])^2 -
    bK * (theta_w2 - theta_w^2 / s1 - (K - 1))

  dval <- bK * denom
  cdx <- cval + dval * xmat

  J <- matrix(0, nrow = len_zi, ncol = len_x)
  ok_J <- cdx >= 0
  J[ok_J] <- abs(denom / (2 * sqrt(cdx[ok_J])))

  g_sqrts <-
    theta_w^2 / (s1 * sig2i[K])^2 -
    bK * (theta_w2 - theta_w^2 / s1 - (denom * xmat + (K - 1)))

  gL <- matrix(0, nrow = len_zi, ncol = len_x)
  gU <- matrix(0, nrow = len_zi, ncol = len_x)

  ok_g <- g_sqrts >= 0

  gL[ok_g] <-
    (theta_w[ok_g] / (s1 * sig2i[K]) - sqrt(g_sqrts[ok_g])) / bK

  gU[ok_g] <-
    (theta_w[ok_g] / (s1 * sig2i[K]) + sqrt(g_sqrts[ok_g])) / bK

  sdK <- sqrt(sig2i[K] + tau2)

  dgL <- stats::dnorm(gL, mean = 0, sd = sdK)
  dgU <- stats::dnorm(gU, mean = 0, sd = sdK)

  fztau2x <- J * (dgL + dgU) / 2
  fztau2x1 <- J * (gL * dgL + gU * dgU) / 2
  fztau2x2 <- J * (gL^2 * dgL + gU^2 * dgU) / 2

  list(
    fztau2 = t(fztau2x),
    g1 = t(fztau2x1),
    g2 = t(fztau2x2)
  )
}

#' Conditional variance of the DL test statistic
#'
#' @keywords internal
aedl_v2tdls <- function(cov_list, x, sig2i, K) {
  x_trunc <- pmax(x, 0)

  vxts <- 1 / (
    matrix(sig2i, nrow = length(x), ncol = K, byrow = TRUE) +
      matrix(x_trunc, nrow = length(x), ncol = K)
  )

  vxt <- rowSums(vxts)
  v_tdl <- numeric(length(x))

  for (ix in seq_along(x)) {
    weight_outer <- tcrossprod(vxts[ix, ])
    v_tdl[ix] <- sum(cov_list[[ix]] * weight_outer) / vxt[ix]
  }

  v_tdl
}
