# Copyright (C) 2021- Ioannis Kosmidis

#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 or 3 of the License
#  (at your option).
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  A copy of the GNU General Public License is available at
#  http://www.r-project.org/Licenses/

#' Ordinal superiority scores of Agresti and Kateri (2017)
#'
#' [ordinal_superiority()] is a method for the estimation and
#' inference about model-based ordinal superiority scores introduced
#' in Agresti and Kateri (2017, Section 5) from fitted objects. The
#' mean bias of the estimates of the ordinal superiority scores can be
#' corrected.
#'
#' @aliases ordinal_superiority
#'
#' @param object a fitted object from an ordinal regression
#'     model. Currently only models from class [`"bracl"`][bracl] are supported.
#' @param formula a RHS formula indicating the group variable to use.
#' @param data an optional data frame in which to look for variables
#'     with which to compute ordinal superiority measures.  If
#'     omitted, an attempt is made to use the data that produced
#'     `object`.
#' @param measure either `"gamma"` (default) or `"Delta"`, specifying
#'     the ordinal superiority measure to be returned.
#' @param level the confidence level required when computing
#'     confidence intervals for the ordinal superiority measures.
#' @param bc logical. If `FALSE` (default) then the ordinal
#'     superiority measures are computed using the estimates in
#'     `object`. If `TRUE` then the ordinal superiority measure
#'     estimates are corrected for mean bias.
#'
#' @references
#'
#' Agresti, A., Kateri, M. (2017). Ordinal probability effect measures
#' for group comparisons in multinomial cumulative link models.
#' *Biometrics*, **73** 214-219. \doi{10.1111/biom.12565}.
#'
#' @examples
#'
#' data("stemcell", package = "brglm2")
#'
#' # Adjacent category logit (proportional odds)
#' stem <- within(stemcell, {nreligion = as.numeric(religion)})
#' fit_bracl_p <- bracl(research ~ nreligion + gender, weights = frequency,
#'                      data = stem, type = "ML", parallel = TRUE)
#'
#' # Estimates and 95% confidence intervals for the probabilities that the response
#' # category for gender "female" is higher than the response category for gender "male",
#' # while adjusting for religion.
#' ordinal_superiority(fit_bracl_p, ~ gender)
#'
#' \donttest{
#' # And their (very-similar in value here) bias corrected versions
#' # with 99% CIs
#' ordinal_superiority(fit_bracl_p, ~ gender, bc = TRUE, level = 0.99)
#' # Note that the object is refitted with type = "AS_mean"
#'
#' }
#'
#' @export
ordinal_superiority.bracl <- function(object, formula, data,
                                      measure = c("gamma", "Delta"),
                                      level = 0.95,
                                      bc = FALSE) {
    measure <- match.arg(measure)
    if (!inherits(object, "bracl")) {
        stop("ordinal superiority measures are not available for objects of class ", class(object)[1])
    }
    ## If bc is TRUE and object is not a reduced mean-bias fit then
    ## compute reduced mean-bias estimators.
    if (isTRUE(bc) & !(object$type %in% c("AS_mean", "AS_mixed"))) {
        object <- update(object, type = "AS_mean")
        warning('`object` was refitted using `type = "AS_mean"`.')
    }
    source_data <- if (missing(data)) {
        if (!is.null(object$call$data)) {
            eval(object$call$data, environment(formula(object)), parent.frame())
        } else {
            environment(formula(object))
        }
    } else {
        data
    }
    source_data <- get_all_vars(delete.response(object$terms), source_data)
    row_ids <- if (missing(data)) {
        row.names(model.frame(object))
    } else {
        row.names(model.frame(delete.response(object$terms), source_data,
                              na.action = na.omit, xlev = object$xlevels))
    }
    source_data <- source_data[row_ids, , drop = FALSE]
    mf <- model.frame(formula, source_data)
    Terms <- attr(mf, "terms")
    z <- model.matrix(Terms, mf, object$contrasts[all.vars(formula)])
    znames <- colnames(z)[-match("(Intercept)", colnames(z), nomatch = 0L)]
    if (isTRUE(length(znames) > 1) | !isTRUE(all(sort(unique(z[, znames])) %in% c(0, 1)))) {
        stop("`formula` can have only one grouping explanatory variable with two levels")
    }
    group_name <- all.vars(formula)
    if (length(group_name) != 1L) {
        stop("`formula` can have only one grouping explanatory variable with two levels")
    }
    predictor_data <- source_data[, !(names(source_data) %in% c(as.character(formula(object)[[2L]]), "(weights)")), drop = FALSE]
    if (!(group_name %in% names(predictor_data))) {
        stop("the grouping explanatory variable must be one of the explanatory variables in `object`")
    }
    group_values <- if (is.factor(mf[[1L]])) levels(droplevels(mf[[1L]])) else sort(unique(mf[[1L]]))
    base_data <- unique(predictor_data[, setdiff(names(predictor_data), group_name), drop = FALSE])
    X0data <- X1data <- base_data
    X0data[[group_name]] <- rep(group_values[1L], nrow(base_data))
    X1data[[group_name]] <- rep(group_values[2L], nrow(base_data))
    X0 <- model.matrix(object, data = X0data)
    X1 <- model.matrix(object, data = X1data)
    same_cols <- colSums(abs(X0 - X1)) == 0
    Xnoz <- X0[, same_cols, drop = FALSE]
    nx <- nrow(X0)
    gammas <- .ordsup(coef(object),
                      X0 = X0, X1 = X1, ncat = object$ncat,
                      ref = object$ref, lev = object$lev, measure = "gamma",
                      po = object$parallel)
    coef_vcov <- vcov(object)
    ## mean bias reduction of gammas
    bias_gammas <- numeric(length(gammas))
    if (bc) {
        for (j in seq.int(nx)) {
            hess <- numDeriv::hessian(function(theta, ...) .ordsup(theta, ...)[j],
                                      x = coef(object),
                                      X0 = X0, X1 = X1, ncat = object$ncat,
                                      ref = object$ref, lev = object$lev, measure = "gamma",
                                      po = object$parallel)
            bias_gammas[j] <- sum(diag(coef_vcov %*% hess)) / 2
        }
    }
    gammas <- gammas - bias_gammas
    ## mean_gammas <- mean(gammas)
    ## compute standard error for gamma
    grads <- numDeriv::jacobian(.ordsup, x = coef(object),
                                X0 = X0, X1 = X1, ncat = object$ncat,
                                ref = object$ref, lev = object$lev, measure = "gamma",
                                po = object$parallel)
    ## grad_mean <- apply(grads, 2, mean)
    se <- apply(grads, 1, function(x) sqrt(crossprod(x, (coef_vcov %*% x))))
    ## se_mean <- sqrt(crossprod(grad_mean, (coef_vcov %*% grad_mean)))
    ## Confidence intervals as in Agresti and Kateri
    a <- 1/2 + level/2
    pct <- paste(format(100 * c(1 - a, a), trim = TRUE, scientific = FALSE, digits = 3), "%")
    lsd <- drop(qnorm(a) * se / (gammas * (1 - gammas)))
    ci <- qlogis(gammas) + cbind(rep(-1, nx),  1) * lsd
    ## lsd_mean <- drop(qnorm(a) * se_mean / (mean_gammas * (1 - mean_gammas)))
    ## ci_mean <- qlogis(mean_gammas) + c(-1, 1) * lsd_mean
    if (isTRUE(measure == "Delta")) {
        out <- cbind(Xnoz, 2 * gammas - 1, 2 * se, 2 * plogis(ci) - 1)
        ## out_mean <- c(2 * mean_gammas - 1, 2 * se_mean, 2 * plogis(ci_mean) - 1)
        colnames(out)[ncol(Xnoz) + 1:4] <- c("Delta", "se", pct)
        ## names(out_mean) <- c("Delta*", "se", pct)
    } else {
        out <- cbind(Xnoz, gammas, se, plogis(ci))
        ## out_mean <- c(mean_gammas, se_mean, plogis(ci_mean))
        colnames(out)[ncol(Xnoz) + 1:4] <- c("gamma", "se", pct)
        ## names(out_mean) <- c("gamma*", "se", pct)
    }
    ## list(individual = out, mean = out_mean)
    out
}

.ordsup <- function(coef, X0, X1, ncat, ref, lev, measure = "gamma", po = TRUE) {
    nams <- names(coef)
    int <- (ncat - 1):1
    sl <- nams[-int]
    if (po) {
        coefs <- cbind(rev(cumsum(coef[int])),
                       int * matrix(coef[sl], nrow = ncat - 1, ncol = length(sl), byrow = TRUE))
    } else {
        coefs <- matrix(coef, nrow = ncat - 1)
        coefs <- apply(coefs, 2, function(x) rev(cumsum(x[int])))
    }
    rownames(coefs) <- lev[-ref]
    probs0 <- .bracl_probs(coefs, X0, ncat, lev)
    probs1 <- .bracl_probs(coefs, X1, ncat, lev)
    gamma_fun <- function(p0, p1) {
        out <- outer(p0, p1, "*")
        sum(out[upper.tri(out)]) + sum(diag(out)) / 2
    }
    out <- numeric(nrow(X0))
    for (i in seq_len(nrow(X0))) {
        out[i] <- gamma_fun(probs0[i, ], probs1[i, ])
    }
    if (measure == "gamma") out else 2 * out - 1
}

.bracl_probs <- function(coefs, X, ncat, lev) {
    fits <- matrix(0, nrow = nrow(X), ncol = ncat, dimnames = list(rownames(X), lev))
    fits1 <- apply(coefs, 1, function(b) X %*% b)
    fits[, rownames(coefs)] <- fits1
    t(apply(fits, 1, function(x) exp(x) / sum(exp(x))))
}
