# Internal input validation
# @keywords internal
.validate_inputs <- function(W, T, G) {
  if (any(is.na(T))){
    stop("Missing data not allowed in T")
  }
  if (any(is.na(W))){
    stop("Missing data not allowed in W")
  }

  if (length(W) != length(T) || length(W) != length(G)) {
    stop("W, T, and G must have the same length.", call. = FALSE)
  }
  if (any(!is.finite(W))) {
    stop("W must be finite.", call. = FALSE)
  }
  if (!all(T %in% c(0, 1))) {
    stop("T must be binary with values 0/1.", call. = FALSE)
  }
  invisible(TRUE)
}

.log1pexp <- function(z) ifelse(z > 35, z, log1p(exp(z)))

.numeric_hessian <- function(fn, x, step = 1e-4) {
  p <- length(x)
  H <- matrix(0, p, p)
  fx <- fn(x)
  for (j in seq_len(p)) {
    ej <- rep(0, p); ej[j] <- step
    f1 <- fn(x + ej)
    f2 <- fn(x - ej)
    H[j, j] <- (f1 - 2 * fx + f2) / step^2
  }
  if (p >= 2) {
    for (j in 1:(p - 1)) {
      for (k in (j + 1):p) {
        ej <- rep(0, p); ej[j] <- step
        ek <- rep(0, p); ek[k] <- step
        fpp <- fn(x + ej + ek)
        fpm <- fn(x + ej - ek)
        fmp <- fn(x - ej + ek)
        fmm <- fn(x - ej - ek)
        H[j, k] <- (fpp - fpm - fmp + fmm) / (4 * step^2)
        H[k, j] <- H[j, k]
      }
    }
  }
  H
}

.c_estep <- function(W, T, G, par, grid_size, grid_sd, delta0_fixed = FALSE) {
  if (delta0_fixed) {
      .Call("c_e_step_loglik_fixed", W, as.integer(T), as.integer(G),
            as.numeric(par), as.integer(grid_size), as.numeric(grid_sd))
  } else {
      .Call("c_e_step_loglik_free", W, as.integer(T), as.integer(G),
            as.numeric(par), as.integer(grid_size), as.numeric(grid_sd))
  }
}

.c_loglik <- function(W, T, G, par, grid_size, grid_sd, delta0_fixed = FALSE) {
  if (delta0_fixed) {
      .Call("c_loglik_fixed", W, as.integer(T), as.integer(G), as.numeric(par),
            as.integer(grid_size), as.numeric(grid_sd))[[1]]
  } else {
      .Call("c_loglik_free", W, as.integer(T), as.integer(G), as.numeric(par),
            as.integer(grid_size), as.numeric(grid_sd))[[1]]
  }
}

.c_q_alpha <- function(xgrid, wpost, T, G, alpha, mu0) {
    .Call("c_q_alpha", xgrid, wpost, as.integer(T), as.integer(G),
          as.numeric(alpha), as.numeric(mu0))[[1]]
}

.make_em_result <- function(estimates, post_mean, post_second, loglik_trace, converged, iterations, se = NULL, vcov = NULL, call = NULL) {
  structure(list(
    estimates = estimates,
    posterior_mean_X = post_mean,
    posterior_second_X = post_second,
    logLik = tail(loglik_trace, 1),
    logLik_trace = loglik_trace,
    converged = converged,
    iterations = iterations,
    se.est = se,
    vcov = vcov,
    tau = estimates[2],
    se = se[2],
    call = call
  ), class = "rctme_fit")
}

.make_stan_result1 <- function(fit, call = NULL) {
    pars <- c("mu0", "tau", "sigmaX", "sigmaW0", "sigmaW1", "sigmaX2",
              "sigmaW0_2", "sigmaW1_2",
              "alpha0", "alphaT", "alphaX", "delta1")
    summ <- as.data.frame(rstan::summary(fit, pars = pars)$summary)

    draws <- as.data.frame(rstan::extract(fit, pars = "tau"))
    taux <- draws$tau
    tau.hat <- mean(taux, trim=0.1)
    
    draws <- as.data.frame(rstan::extract(fit, pars = "alphaX"))
    ax <- draws$alphaX
    ax.mean <- mean(ax)
    ax.ci <- quantile(ax, c(0.025, 0.5, 0.975))
    ROPE.05 <- mean(abs(ax) < 0.05) # example ROPE
    ROPE.15 <- mean(abs(ax) < 0.15) # example ROPE
    ROPE.25 <- mean(abs(ax) < 0.25) # example ROPE
    ROPE.50 <- mean(abs(ax) < 0.50) # example ROPE
    ax.ci1 <- mean(ax > 0)
    ax.out <- list(mean=ax.mean, median=ax.ci[[2]],
                   ci = c(ax.ci[[1]],ax.ci[[3]]),
                   ROPE.05 = ROPE.05,
                   ROPE.15 = ROPE.15,
                   ROPE.25 = ROPE.25,
                   ROPE.50 = ROPE.50,
                   p.gt.zero = ax.ci1)

    structure(list(
        fit = fit,
        summary = summ,
        tau = tau.hat,
        se = summ[2,2],
        contamination = ax.out,
        call = call
    ), class = "rctstan_fit")
}

.make_stan_result2 <- function(fit, call = NULL) {
    pars <- c("mu0", "tau", "sigmaX", "sigmaW0", "sigmaW1", "sigmaX2",
              "sigmaW0_2", "sigmaW1_2", "alpha0", "alphaT", "alphaX",
              "delta0", "delta1")
    summ <- as.data.frame(rstan::summary(fit, pars = pars)$summary)
    ##list(fit = fit, summary = summ, stan_file = stan_file)

    draws <- as.data.frame(rstan::extract(fit, pars = "alphaX"))
    ax <- draws$alphaX
    ax.mean <- mean(ax)
    ax.ci <- quantile(ax, c(0.025, 0.5, 0.975))
    ROPE.05 <- mean(abs(ax) < 0.05) # example ROPE
    ROPE.15 <- mean(abs(ax) < 0.15) # example ROPE
    ROPE.25 <- mean(abs(ax) < 0.25) # example ROPE
    ROPE.50 <- mean(abs(ax) < 0.50) # example ROPE
    ax.ci1 <- mean(ax > 0)
    ax.out <- list(mean=ax.mean, median=ax.ci[[2]],
                   ci = c(ax.ci[[1]],ax.ci[[3]]),
                   ROPE.05 = ROPE.05,
                   ROPE.15 = ROPE.15,
                   ROPE.25 = ROPE.25,
                   ROPE.50 = ROPE.50,
                   p.gt.zero = ax.ci1)
    
    structure(list(
        fit = fit,
        summary = summ,
        tau = summ[2,1],
        se = summ[2,2],
        contamination = ax.out,
        call = call
    ), class = "rctstan_fit")
}

# Simulate a small RCT measurement-error example dataset
#
# Generates a toy dataset for examples and testing under the latent outcome
# model implemented in this package.
#
# @param n Sample size.
# @param mu0 Mean of latent outcome in controls.
# @param tau Treatment effect on latent outcome.
# @param sigmaX SD of latent outcome noise.
# @param alpha0,alphaT,alphaX Logistic-response coefficients.
# @param delta0,delta1 Mean measurement error shifts by response group.
# @param sigmaW0,sigmaW1 Measurement-error SDs by response group.
# @param seed Optional RNG seed.
#
# @return A data frame with columns \code{W}, \code{T}, \code{G}, and latent \code{X}.
# @export
simulate_rctme_data <- function(n = 200,
                                mu0 = 0.5,
                                tau = 1,
                                sigmaX = 1,
                                alpha0 = -0.5,
                                alphaT = 0.4,
                                alphaX = 1,
                                delta0 = 0,
                                delta1 = -0.4,
                                sigmaW0 = 0.8,
                                sigmaW1 = 1,
                                seed = NULL) {
  if (!is.null(seed)) set.seed(seed)
  T <- stats::rbinom(n, 1, 0.5)
  X <- stats::rnorm(n, mu0 + tau * T, sigmaX)
  p <- stats::plogis(alpha0 + alphaT * T + alphaX * (X - mu0))
  G <- stats::rbinom(n, 1, p)
  U <- ifelse(G == 0,
              stats::rnorm(n, delta0, sigmaW0),
              stats::rnorm(n, delta1, sigmaW1))
  W <- X + U
  data.frame(W = W, T = T, G = G, X = X)
}


#' @export
print.rctme_fit <- function(x, ...) {
  cat("MLE with EM\n")
  cat("Converged:", x$converged, "\n")
  cat("Iterations:", x$iterations, "\n")
  cat("Log-likelihood:", format(x$logLik, digits = 6), "\n\n")
  print(round(x$estimates, 4))
  if (!is.null(x$se.est)) {
    cat("\nApproximate SEs:\n")
    print(round(x$se.est, 4))
  }
  invisible(x)
}

#' @export
print.rctstan_fit <- function(x, ...) {
  cat("Bayesian with stan\n")
  print(round(x$summary, 4))
  if (!is.null(x$contamination)) {
      cat("\nContamination:\n\t95% CI=(", round(x$contamination$ci[1],4),
          ", ", round(x$contamination$ci[2],4), ")",
          "\n\tP(alphaX>0)=", round(x$contamination$p.gt.zero,4),
          "\n\tP(|alphaX|<0.05)=", round(x$contamination$ROPE.05,4),
          "\n\tP(|alphaX|<0.15)=", round(x$contamination$ROPE.15,4),
          "\n\tP(|alphaX|<0.25)=", round(x$contamination$ROPE.25,4),
          "\n\tP(|alphaX|<0.50)=", round(x$contamination$ROPE.50,4))
  }
  if (!is.null(x$tau)) {
    cat("\n\nApproximate MLE of tau: ", round(x$tau, 4))
  }
  if (!is.null(x$se)) {
    cat("\nApproximate SE of tau: ", round(x$se, 4),"\n\n")
  }
  invisible(x)
}

#' @export
print.rctcpe_fit <- function(x, ...) {
  cat("Change-point estimate\n")
  print(round(x$summary, 4))
  if (!is.null(x$tau)) {
    cat("\nApproximate MLE of tau: ", round(x$tau, 4),"\n\n")
  }
  invisible(x)
}
