
#' Performs inference on two-phase studies data via semiparametric maximum likelihood.
#'
#'
#' @param formula  regression formula, note that if it does not contain the
#'   missing-by-design variable, \code{miscov}, it will return results under
#'   the null hypothesis. Hypothesis testing corresponds to the score
#'   statistic. Otherwise, estimates and hypothesis testing occur under the
#'   alternative hypothesis leading to Wald statistics. All the elements in
#'   \code{formula} except \code{miscov} must be present in \code{data0} and
#'   \code{data1}.
#' @param miscov right hand side formula with the missing-by-design
#'   covariate(s), i.e. the potential causal locus (loci). Must be present in
#'   \code{data1} but absent in \code{data0}.
#' @param auxvar  right hand side formula with the auxiliary variable(s),
#'   i.e. the GWAS SNP from phase 1. Must be present in \code{data0} and
#'   \code{data1}.
#' @param family member of the exponential family (see \code{\link{family}};
#'   \sQuote{quasi} models not available). Default \code{gaussian()}.
#' @param data a data frame containing data from both phases. Phase 2 data is assumed to be the result of the \code{complete.cases} with both \code{miscov} and \code{auxvar}. Missing values in variables other than \code{miscov} are not allowed and may lead to unexpected errors.
#' @param start.values a named list with initial values for the regression
#'   parameters (\code{betas}) and/or the joint distribution between
#'   \code{miscov} and \code{auxvar} (\code{q}). Defaults to \code{NULL}.
#' @param control a list of parameters for controlling the fitting process, i.e., number or EM iterations, tolerance and whether a trace is requested, it is passed to \code{\link{twoPhaseSPML.control}}.
#'
#' @details
#' The function implements the semiparametric maximum likelihood (SPML)
#' estimator for two-phase genetic association studies described by
#' Espin-Garcia et al. The EM algorithm alternates between an E-step that
#' computes observation weights based on the current estimate of the joint
#' distribution of the missing covariate and the auxiliary variable, and an
#' M-step that updates the regression coefficients and the non-parametric
#' distribution via iteratively re-weighted least squares and a closed-form
#' update, respectively.
#'
#' When \code{miscov} is absent from \code{formula}, inference is performed
#' under \eqn{H_0} and score statistics are returned. When \code{miscov} is
#' included in \code{formula}, inference is performed under \eqn{H_a} and
#' Wald statistics are returned.
#'
#' @return An object of class \code{"twoPhaseSPML"}, which is a list
#'   containing:
#'   \describe{
#'     \item{theta}{named numeric vector of estimated regression coefficients.}
#'     \item{qGZ}{data frame with the estimated joint distribution of
#'       \code{miscov} and \code{auxvar}.}
#'     \item{dispersion}{estimated dispersion parameter (variance for
#'       Gaussian; 1 for other families).}
#'     \item{iter}{integer, number of EM iterations performed.}
#'     \item{ll}{log-likelihood at convergence.}
#'     \item{df}{degrees of freedom for the test statistic.}
#'     \item{Sobs}{(null model only) observed score test statistic.}
#'     \item{Sexp}{(null model only) expected score test statistic.}
#'     \item{var_theta}{(alternative model only) named numeric vector of
#'       marginal variances for \code{theta}.}
#'     \item{Wobs}{(alternative model only) observed Wald test statistic.}
#'     \item{Wexp}{(alternative model only) expected Wald test statistic.}
#'     \item{qG}{data frame with the estimated marginal distribution of
#'       \code{miscov}.}
#'     \item{formula}{the \code{formula} used.}
#'     \item{miscov}{the \code{miscov} formula used.}
#'     \item{family}{the \code{family} object used.}
#'     \item{null.model}{logical, \code{TRUE} when the model was fitted under
#'       \eqn{H_0}.}
#'   }
#'
#'
#' @examples
#'
#' data.table::setDTthreads(1)
#' data <- DataGeneration_TPD(N=200)
#' set.seed(1)
#' R <- rep(0, nrow(data))
#' R[sample(nrow(data), 175)] <- 1   # random phase 2 subsample of 500
#' data[R == 0, c("G1")] <- NA
#'
#' # Null model (score test)
#' res_Ho <- twoPhaseSPML(
#'   formula = Y ~ Z,
#'   miscov  = ~ G1,
#'   auxvar  = ~ Z,
#'   data   = data
#' )
#' print(res_Ho)
#'
#' # Alternative model (Wald test)
#' res_Ha <- twoPhaseSPML(
#'   formula = Y ~ Z + G1,
#'   miscov  = ~ G1,
#'   auxvar  = ~ Z,
#'   data   = data
#' )
#' print(res_Ha)
#'
#' @export
#' @import stats
#' @import data.table
#' @importFrom MASS ginv
#' @importFrom Matrix fac2sparse t crossprod colSums
twoPhaseSPML <- function(formula, miscov, auxvar, family = gaussian, data, start.values = NULL, control=list())
{
  if (is.character(family))
    family <- get(family, mode = "function", envir = parent.frame())
  if (is.function(family))
    family <- family()
  if (is.null(family$family)) {
    message(deparse(family))
    stop("'family' not recognized")
  }

  terms_dat <- unique(c(all.vars(formula),
                         attr(terms(miscov), "term.labels"),
                         attr(terms(auxvar), "term.labels")))

  Y_ <- as.character(formula[[2]])
  G_  <- all.vars(miscov)
  Z_  <- all.vars(auxvar)
  terms_data0 <- if (all(G_ %in% all.vars(formula))) {
    c(all.vars(formula)[-which(all.vars(formula) %in% G_)], Z_)
  } else {
    c(all.vars(formula), Z_)
  }
  terms_data0 <- unique(terms_data0)
  
  ### construct data0 and data1 from data, where data0 is a data frame with the complement of the phase 2 data, must contain all variables in \code{formula} and \code{auxvar} but \emph{not} \code{miscov}; data1 is a data frame with the phase 2 data, must contain all variables
  Ridx <- as.numeric(complete.cases(data[c(G_,Z_)]))
  data0 <- data[Ridx==0, terms_data0]
  data1 <- data[Ridx==1, ]
  
  n1 <- nrow(data1); n0 <- nrow(data0); n <- n1 + n0
  nG_ <- length(G_)
  uniqZ <- unique(data1[, Z_]); uniqZ <- uniqZ[order(uniqZ)]

  uniqG <- unique(data1[, G_])
  dat0 <- cbind(id_ = rep(seq(n1 + 1, n, by = 1), each = NROW(uniqG)),
                data0[rep(seq_len(n0), each = NROW(uniqG)), ],
                matrix(rep(t(uniqG), n0), ncol = NCOL(uniqG), byrow = TRUE))
  names(dat0)[(ncol(dat0) - NCOL(uniqG) + 1):ncol(dat0)] <- G_
  qform <- as.formula(paste0("wg_ ~", paste(G_, collapse = "+"), "+", Z_))

  dat <- rbind(cbind(id_ = 1:n1, data1[, terms_dat]),
               dat0[, c("id_", terms_dat)])

  # Initialise the non-parametric terms: p_{g,z}
  uniqGZ <- unique(data1[, c(G_, Z_)])
  q0 <- cbind(uniqGZ, q = 1 / nrow(uniqGZ))
  Zuniq <- unique(c(data1[, Z_], data0[, Z_]))
  q <- cbind(matrix(rep(t(uniqG), length(Zuniq)),
                    ncol = NCOL(uniqG), byrow = TRUE),
             Z = rep(Zuniq, each = NROW(uniqG)))
  colnames(q) <- c(G_, Z_)
  q <- merge(q, q0, all.x = TRUE, sort = FALSE)
  q[is.na(q$q), "q"] <- 0
  nq <- nrow(q)
  if (!is.null(start.values$q)) {
    q1 <- merge(q, start.values$q, by = c(G_, Z_), all.x = TRUE, sort = FALSE)
    q1 <- q1[, c(G_, Z_, "q.y")]
    names(q1)[3] <- "q"
    q <- q1
    # normalization check
    total_q <- sum(q[, "q"], na.rm = TRUE)
    if (abs(total_q - 1) > 1e-8)
      warning("start.values$q does not sum to 1 (sum = ", total_q,"); normalizing.")
    q[, "q"] <- q[, "q"] / total_q
  }

  # Initialise regression parameters
  X <- model.matrix(formula, dat)

  if (is.null(start.values$betas)) {
    theta.o <- glm(formula, data = data1, family = family)$coef
    if(any(is.na(theta.o))) theta.o <- glm(formula, data = dat, family = family)$coef
  } else {
    theta.o <- start.values$betas
  }

  nparms <- length(theta.o)
  
  if (family$family %in% c("gaussian", "Gamma")) {
    # disp <- var(c(data1[, Y_], data0[, Y_]))
    glm_init <- glm(formula, data = dat, family = family)
    theta.o   <- glm_init$coef
    disp      <- sum(glm_init$weights * glm_init$residuals^2) / (n - nparms)
  } else {
    disp <- 1
  }

  ng <- NROW(uniqG)
  n1ones <- rep(1, n1)

  # EM convergence settings
  control <- do.call("twoPhaseSPML.control", control)
  maxiter <- control$EMmaxit
  tol     <- control$EMtol
  iter    <- 0L
  theta   <- theta.o
  theta.o <- theta + 1
  q.o     <- q
  q.o[, "q"] <- q.o[, "q"] - 1
  if (control$trace) message("Initial iteration:", iter, "theta =", theta,
                   "pGZ =", as.numeric(q[, "q"]), "\n")

  while ( iter <= maxiter &
         (max(abs(theta - theta.o)) > tol |
          max(abs(q[, "q"] - q.o[, "q"])) > tol) ) {

    theta.o <- theta; q.o <- q; iter <- iter + 1

    # E-step
    dat$wg_ <- c(n1ones,
                 fnc_wgts(formula, family, dat0, theta, q, disp, G_, Z_, uniqZ))

    # M-step
    glmfit <- suppressWarnings(
      irls_svdnewton(X, dat[, Y_], family = family,
                     weights = dat$wg_, rank_deficiency = "select columns")
    )
    theta <- glmfit$coef

    if (family$family %in% c("gaussian", "Gamma")) {
      # disp <- sum(dat$wg_ * glmfit$residuals^2) / (n - nparms)
      mu_fit  <- family$linkinv(X %*% theta)
      pearson <- (dat[, Y_] - mu_fit) / sqrt(family$variance(mu_fit))
      disp    <- sum(dat$wg_ * pearson^2) / (n - nparms)
    }
    q.it <- data.frame(xtabs(qform, data = dat) / n)

    ordq <- match(apply(q[, c(G_, Z_)], 1, paste, collapse = "-"),
                  apply(q.it[, c(G_, Z_)], 1, paste, collapse = "-"))
    q[, "q"] <- q.it[ordq, "Freq"]

    if (control$trace) message("Iteration:", iter, "theta =", theta,
                     "pGZ =", as.numeric(q[, "q"]), "\n")
  }
  # Final wg calculation at convergence (or in case a dry run is requested)
  dat$wg_ <- c(n1ones,
              fnc_wgts(formula, family, dat0, theta, q, disp, G_, Z_, uniqZ))

  if (control$trace) message("Total iterations:", iter,
                   "theta_end =", theta,
                   "pGZ_end =", as.numeric(q[, "q"]), "\n")

  # Objects for hypothesis testing
  qG <- aggregate(as.formula(paste0("q~", paste(G_, collapse = "+"))), q, FUN = sum)
  qG$k_ <- seq_len(nrow(qG))
  dat$ord_ <- seq_len(nrow(dat))
  dat <- merge(dat, qG[, c(G_, "k_")], by = G_)
  dat <- dat[order(dat$ord_), ]

  Ho <- !all(G_ %in% names(theta))   # TRUE when fitting under H0

  if (Ho) {
    indx_thetas_H0 <- which(!G_ %in% names(theta))
    nG_ <- length(indx_thetas_H0)
    val <- c(theta, rep(0, nG_))
    id  <- c(seq_along(theta), 1 + seq(0.1, 0.9, length.out = nG_))
    theta1 <- val[order(id)]
    names(theta1)[which(names(theta1) == "")] <- G_[indx_thetas_H0]
    formula1 <- update(formula,
                       paste0(".~", paste(G_[indx_thetas_H0], collapse = "+"), "+ ."))
    X1  <- model.matrix(formula1, dat)
    eta1 <- X1 %*% theta1
    W   <- family$mu.eta(eta1)^2 / family$variance(family$linkinv(eta1))
    glmfit1 <- list()
    glmfit1$qr <- qr(X1 * sqrt(as.numeric(dat$wg_ * W)))

    FIM    <- FIM.glm(theta1, qG, disp, formula1, Y_, dat, family, glmfit1)
    expFIM <- expFIM.glm(theta1, qG, disp, formula1, Y_, dat, family, glmfit1)

    res_SobsIM <- Testing_EM_joint(theta1, qG, NULL, G_[indx_thetas_H0], FIM)
    res_SexpIM <- Testing_EM_joint(theta1, qG, NULL, G_[indx_thetas_H0], expFIM)

  } else {
    FIM    <- FIM.glm(theta, qG, disp, formula, Y_, dat, family, glmfit)
    expFIM <- expFIM.glm(theta, qG, disp, formula, Y_, dat, family, glmfit)

    res_VobsIM <- Testing_EM_joint(theta, qG, NULL, G_, FIM)
    res_VexpIM <- Testing_EM_joint(theta, qG, NULL, G_, expFIM)
  }

  ll <- .loglik(theta, q, disp, formula, Y_, G_, Z_, dat, family)

  if (Ho) {
    out <- list(theta      = theta,
                qGZ        = q,
                dispersion = disp,
                Sobs       = res_SobsIM$S,
                Sexp       = res_SexpIM$S,
                df         = res_SexpIM$df,
                qG         = qG,
                iter       = as.integer(iter),
                ll         = ll,
                formula    = formula,
                miscov     = miscov,
                auxvar     = auxvar,
                family     = family,
                null.model = TRUE)
  } else {
    out <- list(theta      = theta,
                qGZ        = q,
                dispersion = disp,
                var_theta  = res_VobsIM$Var,
                Wobs       = res_VobsIM$W,
                Wexp       = res_VexpIM$W,
                df         = res_VexpIM$df,
                qG         = qG,
                iter       = as.integer(iter),
                ll         = ll,
                formula    = formula,
                miscov     = miscov,
                auxvar     = auxvar,
                family     = family,
                null.model = FALSE)
  }

  class(out) <- "twoPhaseSPML"
  out
}


#' Print method for \code{twoPhaseSPML} objects
#'
#' Prints a concise, \code{glm}-style summary of a fitted
#' \code{\link{twoPhaseSPML}} model.
#'
#' @param x an object of class \code{"twoPhaseSPML"}, as returned by
#'   \code{\link{twoPhaseSPML}}.
#' @param digits integer; minimum number of significant digits to print.
#'   Passed to \code{\link{printCoefmat}}. Defaults to
#'   \code{max(3L, getOption("digits") - 3L)}.
#' @param \dots further arguments passed to or from other methods (currently
#'   unused).
#'
#' @details
#' Under the \emph{null} model (\code{miscov} absent from \code{formula}),
#' score statistics (\code{Sobs} and \code{Sexp}) are displayed. Under the
#' \emph{alternative} model, Wald statistics (\code{Wobs} and \code{Wexp})
#' together with standard errors and \eqn{p}-values for the missing-by-design
#' covariate(s) are displayed.  The coefficient table always shows the
#' estimated regression coefficients; standard errors and \eqn{p}-values are
#' shown only when \code{var_theta} is available (alternative model).
#'
#' @return \code{x} is returned invisibly.
#'
#' @seealso \code{\link{twoPhaseSPML}}
#'
#' @examples
#' data.table::setDTthreads(1)
#' data <- DataGeneration_TPD(N=200)
#' set.seed(1)
#' R <- rep(0, nrow(data))
#' R[sample(nrow(data), 175)] <- 1
#' data[R == 0, c("G1")] <- NA
#'
#' res_Ho <- twoPhaseSPML(
#'   formula = Y ~ Z,
#'   miscov  = ~ G1,
#'   auxvar  = ~ Z,
#'   data   = data
#' )
#' print(res_Ho)
#'
#' @export
print.twoPhaseSPML <- function(x, digits = max(3L, getOption("digits") - 3L), ...)
{
  cat("\nCall:\n",
      paste0("twoPhaseSPML(formula = ", deparse(x$formula),
             ", miscov = ", deparse(x$miscov),
             ", auxvar = ", deparse(x$auxvar),
             ", family = ", x$family$family, ")"),
      "\n\n", sep = "")

  # ---- Coefficients table -------------------------
  theta <- x$theta
  G_    <- all.vars(x$miscov)

  if (!x$null.model && !anyNA(x$var_theta)) {
    # Alternative model: show Estimate / SE / z / p
    se    <- sqrt(x$var_theta)
    zval  <- theta / se
    pval  <- 2 * pnorm(abs(zval), lower.tail = FALSE)
    coef_mat <- cbind(Estimate    = theta,
                      `Std. Error` = se,
                      `z value`    = zval,
                      `Pr(>|z|)`   = pval)
    cat("Coefficients:\n")
    printCoefmat(coef_mat, digits = digits, signif.stars = TRUE,
                 na.print = "NA", ...)
  } else {
    # Null model or unavailable SE: show Estimate only
    coef_mat <- cbind(Estimate = theta)
    cat("Coefficients:\n")
    printCoefmat(coef_mat, digits = digits, na.print = "NA", ...)
  }

  # ---- Family / link --------------------
  cat("\nFamily:", x$family$family,
      " Link:", x$family$link, "\n")

  # ---- Dispersion -----------------------
  if (x$family$family %in% c("gaussian", "Gamma")) {
    cat("Dispersion:", format(x$dispersion, digits = digits), "\n")
  }

  # ---- Log-likelihood --------------------
  cat("Log-likelihood:", format(x$ll, digits = digits),
      " (EM iterations:", x$iter, ")\n")

  # ---- Test statistics -------------------
  cat("\nHypothesis test for missing-by-design covariate(s):",
      paste(G_, collapse = ", "), "\n")

  if (x$null.model) {
    cat("  [Null model - Score statistics]\n")
    cat("  Observed score statistic (Sobs):", format(x$Sobs, digits = digits), "\n")
    cat("  Expected score statistic (Sexp):", format(x$Sexp, digits = digits), "\n")
  } else {
    cat("  [Alternative model - Wald statistics]\n")
    cat("  Observed Wald statistic (Wobs):", format(x$Wobs, digits = digits), "\n")
    cat("  Expected Wald statistic (Wexp):", format(x$Wexp, digits = digits), "\n")
  }
  cat("  Degrees of freedom:", x$df, "\n")

  if (!is.null(x$df) && !is.na(x$df)) {
    stat <- if (x$null.model) x$Sobs else x$Wobs
    if (!is.null(stat) && !is.na(stat)) {
      pv <- pchisq(stat, df = x$df, lower.tail = FALSE)
      cat("  p-value (chi-squared):", format.pval(pv, digits = digits), "\n")
    }
  }

  cat("\n")
  invisible(x)
}


#' Auxiliary for Controlling the EM algorithm in twoPhaseSPML.
#'
#'
#' @param EMmaxit integer giving the maximal number of iterations of the EM algorithm.
#' @param EMtol numeric positive convergence tolerance for EM algorithm covnergence.
#' @param trace logical \code{TRUE}, prints progress at each EM iteration. Defaults to \code{FALSE}.
#'
#' @details
#' The function implements the semiparametric maximum likelihood (SPML)
#' 
#'
#'
#' @return A list with components named as the arguments.
#' 
#' @export
twoPhaseSPML.control <- function (EMmaxit = 1000, EMtol = 1e-05, trace = FALSE) 
{
  if (!is.numeric(EMtol) || EMtol <= 0) 
    stop("value of 'epsilon' must be > 0")
  if (!is.numeric(EMmaxit)) 
    stop("maximum number of iterations must be numeric")
  list(EMtol = EMtol, EMmaxit = EMmaxit, trace = trace)
}