# EM fit with free `delta0` and `delta1`
#
# Fits the latent-outcome model
# \deqn{X|T \sim N(\mu_0 + \tau T, \sigma_X^2),}
# \deqn{G|X,T \sim \text{Bernoulli}(\text{logit}^{-1}(\alpha_0 + \alpha_T T + \alpha_X (X-\mu_0))),}
# \deqn{W|X,G=g \sim N(X + \delta_g, \sigma_{Wg}^2).}
#
# The E-step and observed-data log-likelihood are evaluated in compiled C code
# for speed.
#
# @param W Numeric surrogate outcome.
# @param T Binary treatment indicator.
# @param G Binary response indicator.
# @param max_iter Maximum EM iterations.
# @param tol Convergence tolerance.
# @param grid_size Number of E-step grid points.
# @param grid_sd Grid half-width in posterior SD units.
# @param verbose Print iteration trace.
# @param compute_se Compute approximate SEs from a numerical observed Hessian.
# @param hess_step Finite-difference step for Hessian approximation.
# @param init Optional named list of initial values.
#
# @return An object of class `rctme_fit`.
# @export
fit_rctme_em <- function(W, T, G,
                         max_iter = 200,
                         tol = 1e-6,
                         grid_size = 81,
                         grid_sd = 6,
                         verbose = FALSE,
                         compute_se = TRUE,
                         hess_step = 1e-4,
                         init = NULL) {
  .validate_inputs(W, T, G)
  W <- as.numeric(W)
  T <- as.integer(T)
  G <- as.integer(G)
  var_floor <- 1e-8

  if (is.null(init)) {
    mu0 <- mean(W[T == 0])
    tau <- mean(W[T == 1]) - mean(W[T == 0])
    sigmaX2 <- stats::var(W)
    if (!is.finite(sigmaX2) || sigmaX2 <= 0) sigmaX2 <- 1
    delta0 <- 0
    delta1 <- 0
    sigmaW0_2 <- 0.25 * sigmaX2
    sigmaW1_2 <- 0.25 * sigmaX2
    fit0 <- try(stats::glm(G ~ T + W, family = stats::binomial()), silent = TRUE)
    if (inherits(fit0, "try-error")) {
      alpha0 <- stats::qlogis(mean(G)); alphaT <- 0; alphaX <- 0
    } else {
      cf <- stats::coef(fit0)
      alpha0 <- ifelse(is.finite(cf[1]), cf[1], stats::qlogis(mean(G)))
      alphaT <- ifelse(is.finite(cf[2]), cf[2], 0)
      alphaX <- ifelse(is.finite(cf[3]), cf[3], 0)
    }
  } else {
    mu0 <- if (!is.null(init$mu0)) init$mu0 else init$beta0; tau <- init$tau; sigmaX2 <- init$sigmaX2
    alpha0 <- init$alpha0; alphaT <- init$alphaT; alphaX <- init$alphaX
    delta0 <- init$delta0; delta1 <- init$delta1
    sigmaW0_2 <- init$sigmaW0_2; sigmaW1_2 <- init$sigmaW1_2
  }

  loglik_trace <- numeric(max_iter)
  post_mean <- post_second <- NULL
  xgrid <- wpost <- NULL
  converged <- FALSE

  for (iter in seq_len(max_iter)) {
    old <- c(mu0, tau, sigmaX2, alpha0, alphaT, alphaX, delta0, delta1, sigmaW0_2, sigmaW1_2)
    sigmaX2 <- max(sigmaX2, var_floor)
    sigmaW0_2 <- max(sigmaW0_2, var_floor)
    sigmaW1_2 <- max(sigmaW1_2, var_floor)

    par <- c(mu0, tau, sigmaX2, alpha0, alphaT, alphaX, delta0, delta1, sigmaW0_2, sigmaW1_2)
    estep <- .c_estep(W, T, G, par, grid_size, grid_sd, delta0_fixed = FALSE)
    post_mean <- estep$post_mean
    post_second <- estep$post_second
    xgrid <- estep$xgrid
    wpost <- estep$wpost
    loglik_trace[iter] <- estep$loglik

    Z <- cbind(1, T)
    beta_hat <- solve(crossprod(Z), crossprod(Z, post_mean))
    mu0 <- unname(beta_hat[1]); tau <- unname(beta_hat[2])
    mu <- mu0 + tau * T
    sigmaX2 <- mean(post_second - 2 * mu * post_mean + mu^2)
    sigmaX2 <- max(sigmaX2, var_floor)

    idx0 <- which(G == 0L)
    idx1 <- which(G == 1L)
    if (length(idx0)) {
      delta0 <- mean(W[idx0] - post_mean[idx0])
      sigmaW0_2 <- mean((W[idx0] - delta0)^2 - 2 * (W[idx0] - delta0) * post_mean[idx0] + post_second[idx0])
      sigmaW0_2 <- max(sigmaW0_2, var_floor)
    }
    if (length(idx1)) {
      delta1 <- mean(W[idx1] - post_mean[idx1])
      sigmaW1_2 <- mean((W[idx1] - delta1)^2 - 2 * (W[idx1] - delta1) * post_mean[idx1] + post_second[idx1])
      sigmaW1_2 <- max(sigmaW1_2, var_floor)
    }

    opt <- stats::optim(c(alpha0, alphaT, alphaX),
      fn = function(a) -.c_q_alpha(xgrid, wpost, T, G, a, mu0),
      method = "BFGS", control = list(maxit = 200, reltol = 1e-10)
    )
    alpha0 <- opt$par[1]; alphaT <- opt$par[2]; alphaX <- opt$par[3]

    new <- c(mu0, tau, sigmaX2, alpha0, alphaT, alphaX, delta0, delta1, sigmaW0_2, sigmaW1_2)
    dpar <- max(abs(new - old))
    dll <- if (iter == 1) Inf else loglik_trace[iter] - loglik_trace[iter - 1]
    if (verbose) message(sprintf("iter=%d logLik=%.6f max|dpar|=%.3e dLL=%.3e", iter, loglik_trace[iter], dpar, dll))
    if (iter > 1 && abs(dll) < tol && dpar < sqrt(tol)) {
      converged <- TRUE
      loglik_trace <- loglik_trace[seq_len(iter)]
      break
    }
    if (iter == max_iter) loglik_trace <- loglik_trace[seq_len(iter)]
  }

  est <- c(mu0 = mu0, tau = tau, sigmaX2 = sigmaX2,
           alpha0 = alpha0, alphaT = alphaT, alphaX = alphaX,
           delta0 = delta0, delta1 = delta1,
           sigmaW0_2 = sigmaW0_2, sigmaW1_2 = sigmaW1_2)

  se <- NULL; vcov <- NULL
  if (compute_se) {
    theta_hat <- c(mu0, tau, log(sigmaX2), alpha0, alphaT, alphaX, delta0, delta1, log(sigmaW0_2), log(sigmaW1_2))
    obj_theta <- function(theta) {
      par_nat <- c(theta[1], theta[2], exp(theta[3]), theta[4], theta[5], theta[6], theta[7], theta[8], exp(theta[9]), exp(theta[10]))
      .c_loglik(W, T, G, par_nat, grid_size, grid_sd, delta0_fixed = FALSE)
    }
    H <- .numeric_hessian(obj_theta, theta_hat, step = hess_step)
    info <- -H
    Vtheta <- try(solve(info), silent = TRUE)
    if (!inherits(Vtheta, "try-error") && all(is.finite(Vtheta))) {
      J <- diag(c(1, 1, sigmaX2, 1, 1, 1, 1, 1, sigmaW0_2, sigmaW1_2))
      vcov <- J %*% Vtheta %*% t(J)
      colnames(vcov) <- rownames(vcov) <- names(est)
      se <- sqrt(pmax(diag(vcov), 0))
    }
  }
  .make_em_result(est, post_mean, post_second, loglik_trace, converged, length(loglik_trace), se = se, vcov = vcov, call = match.call())
}

# EM fit with `delta0` fixed at 0
#
# Same model as [fit_rctme_em()] but with the measurement-error mean shift in
# the `G = 0` group fixed at zero.
#
# @inheritParams fit_rctme_em
# @return An object of class `rctme_fit`.
# @export
fit_rctme_em_delta0_fixed <- function(W, T, G,
                                      max_iter = 200,
                                      tol = 1e-6,
                                      grid_size = 81,
                                      grid_sd = 6,
                                      verbose = FALSE,
                                      compute_se = TRUE,
                                      hess_step = 1e-4,
                                      init = NULL) {
  .validate_inputs(W, T, G)
  W <- as.numeric(W)
  T <- as.integer(T)
  G <- as.integer(G)
  var_floor <- 1e-8

  if (is.null(init)) {
    mu0 <- mean(W[T == 0])
    tau <- mean(W[T == 1]) - mean(W[T == 0])
    sigmaX2 <- stats::var(W)
    if (!is.finite(sigmaX2) || sigmaX2 <= 0) sigmaX2 <- 1
    delta1 <- 0
    sigmaW0_2 <- 0.25 * sigmaX2
    sigmaW1_2 <- 0.25 * sigmaX2
    fit0 <- try(stats::glm(G ~ T + W, family = stats::binomial()), silent = TRUE)
    if (inherits(fit0, "try-error")) {
      alpha0 <- stats::qlogis(mean(G)); alphaT <- 0; alphaX <- 0
    } else {
      cf <- stats::coef(fit0)
      alpha0 <- ifelse(is.finite(cf[1]), cf[1], stats::qlogis(mean(G)))
      alphaT <- ifelse(is.finite(cf[2]), cf[2], 0)
      alphaX <- ifelse(is.finite(cf[3]), cf[3], 0)
    }
  } else {
    mu0 <- if (!is.null(init$mu0)) init$mu0 else init$beta0; tau <- init$tau; sigmaX2 <- init$sigmaX2
    alpha0 <- init$alpha0; alphaT <- init$alphaT; alphaX <- init$alphaX
    delta1 <- init$delta1; sigmaW0_2 <- init$sigmaW0_2; sigmaW1_2 <- init$sigmaW1_2
  }

  loglik_trace <- numeric(max_iter)
  post_mean <- post_second <- NULL
  xgrid <- wpost <- NULL
  converged <- FALSE

  for (iter in seq_len(max_iter)) {
    old <- c(mu0, tau, sigmaX2, alpha0, alphaT, alphaX, delta1, sigmaW0_2, sigmaW1_2)
    sigmaX2 <- max(sigmaX2, var_floor)
    sigmaW0_2 <- max(sigmaW0_2, var_floor)
    sigmaW1_2 <- max(sigmaW1_2, var_floor)

    par <- c(mu0, tau, sigmaX2, alpha0, alphaT, alphaX, delta1, sigmaW0_2, sigmaW1_2)
    estep <- .c_estep(W, T, G, par, grid_size, grid_sd, delta0_fixed = TRUE)
    post_mean <- estep$post_mean
    post_second <- estep$post_second
    xgrid <- estep$xgrid
    wpost <- estep$wpost
    loglik_trace[iter] <- estep$loglik

    Z <- cbind(1, T)
    beta_hat <- solve(crossprod(Z), crossprod(Z, post_mean))
    mu0 <- unname(beta_hat[1]); tau <- unname(beta_hat[2])
    mu <- mu0 + tau * T
    sigmaX2 <- mean(post_second - 2 * mu * post_mean + mu^2)
    sigmaX2 <- max(sigmaX2, var_floor)

    idx0 <- which(G == 0L)
    idx1 <- which(G == 1L)
    if (length(idx0)) {
      sigmaW0_2 <- mean(W[idx0]^2 - 2 * W[idx0] * post_mean[idx0] + post_second[idx0])
      sigmaW0_2 <- max(sigmaW0_2, var_floor)
    }
    if (length(idx1)) {
      delta1 <- mean(W[idx1] - post_mean[idx1])
      sigmaW1_2 <- mean((W[idx1] - delta1)^2 - 2 * (W[idx1] - delta1) * post_mean[idx1] + post_second[idx1])
      sigmaW1_2 <- max(sigmaW1_2, var_floor)
    }

    opt <- stats::optim(c(alpha0, alphaT, alphaX),
      fn = function(a) -.c_q_alpha(xgrid, wpost, T, G, a, mu0),
      method = "BFGS", control = list(maxit = 200, reltol = 1e-10)
    )
    alpha0 <- opt$par[1]; alphaT <- opt$par[2]; alphaX <- opt$par[3]

    new <- c(mu0, tau, sigmaX2, alpha0, alphaT, alphaX, delta1, sigmaW0_2, sigmaW1_2)
    dpar <- max(abs(new - old))
    dll <- if (iter == 1) Inf else loglik_trace[iter] - loglik_trace[iter - 1]
    if (verbose) message(sprintf("iter=%d logLik=%.6f max|dpar|=%.3e dLL=%.3e", iter, loglik_trace[iter], dpar, dll))
    if (iter > 1 && abs(dll) < tol && dpar < sqrt(tol)) {
      converged <- TRUE
      loglik_trace <- loglik_trace[seq_len(iter)]
      break
    }
    if (iter == max_iter) loglik_trace <- loglik_trace[seq_len(iter)]
  }

  est <- c(mu0 = mu0, tau = tau, sigmaX2 = sigmaX2,
           alpha0 = alpha0, alphaT = alphaT, alphaX = alphaX,
           delta0 = 0, delta1 = delta1,
           sigmaW0_2 = sigmaW0_2, sigmaW1_2 = sigmaW1_2)

  se <- NULL; vcov <- NULL
  if (compute_se) {
    theta_hat <- c(mu0, tau, log(sigmaX2), alpha0, alphaT, alphaX, delta1, log(sigmaW0_2), log(sigmaW1_2))
    obj_theta <- function(theta) {
      par_nat <- c(theta[1], theta[2], exp(theta[3]), theta[4], theta[5], theta[6], theta[7], exp(theta[8]), exp(theta[9]))
      .c_loglik(W, T, G, par_nat, grid_size, grid_sd, delta0_fixed = TRUE)
    }
    H <- .numeric_hessian(obj_theta, theta_hat, step = hess_step)
    info <- -H
    Vtheta <- try(solve(info), silent = TRUE)
    if (!inherits(Vtheta, "try-error") && all(is.finite(Vtheta))) {
      J <- diag(c(1, 1, sigmaX2, 1, 1, 1, 1, sigmaW0_2, sigmaW1_2))
      vcov_free <- J %*% Vtheta %*% t(J)
      free_names <- c("mu0", "tau", "sigmaX2", "alpha0", "alphaT", "alphaX", "delta1", "sigmaW0_2", "sigmaW1_2")
      colnames(vcov_free) <- rownames(vcov_free) <- free_names
      vcov <- matrix(0, 10, 10, dimnames = list(names(est), names(est)))
      idx <- c(1, 2, 3, 4, 5, 6, 8, 9, 10)
      vcov[idx, idx] <- vcov_free
      se <- sqrt(pmax(diag(vcov), 0))
    }
  }
  .make_em_result(est, post_mean, post_second, loglik_trace, converged, length(loglik_trace), se = se, vcov = vcov, call = match.call())
}
