
#' @name residuals
#' @aliases residuals,MSAassess-method
#'
#' @title Calculate model residuals
#'
#' @description Extract residuals from fitted model
#'
#' @param object [MSAassess-class] object returned by `fit_MSA()`
#' @param vars Character vector to indicate which residuals will be calculated.
#' Available choices from [MSAdata-class] object are:
#' "Cinit_mfr", "Cobs_ymfr", "CAAobs_ymafr", "CALobs_ymlfr",
#' "Iobs_ymi", "IAAobs_ymai", "IALobs_ymli", "SC_ymafrs"
#' @param type Character, 'response' for the `log(observed/predicted)` values or 'pearson' for calculating Z-scores.
#' Composition data always use 'pearson'.
#' @param ... Not used
#'
#' @return A named list based on `vars` argument.
#' @aliases residuals
#' @importFrom stats residuals
#' @export
setMethod("residuals", signature(object = "MSAassess"),
          function(object, vars, type = c("response", "pearson"), ...) {

  vars_choices <- c("Cinit_mfr", "Cobs_ymfr", "CAAobs_ymafr", "CALobs_ymlfr",
                    "Iobs_ymi", "IAAobs_ymai", "IALobs_ymli", "SC_ymafrs", "tag_ymarrs")
  if (missing(vars)) vars <- vars_choices
  vars <- match.arg(vars, choices = vars_choices, several.ok = TRUE)
  type <- match.arg(type)
  dat <- get_MSAdata(object)

  res <- list()

  # Init catch ----
  if (any(vars == "Cinit_mfr") && !is.null(object@report$initCB_mfrs)) {
    Cinitres <- log(dat@Dfishery@Cinit_mfr/apply(object@report$initCB_mfrs, 1:3, sum))
    if (type == "pearson") {
      Cinitres <- Cinitres/0.01
    }
    res$Cinit_mfr <- Cinitres
  }

  # Catch ----
  if (any(vars == "Cobs_ymfr")) {
    Cres <- log(dat@Dfishery@Cobs_ymfr/apply(object@report$CB_ymfrs, 1:4, sum))
    if (type == "pearson") {
      if (length(dat@Dfishery@Csd_ymfr)) {
        Csd <- dat@Dfishery@Csd_ymfr
      } else {
        Csd <- 0.01
      }
      Cres <- Cres/Csd
    }
    res$Cobs_ymfr <- Cres
  }


  # Catch at age ----
  if (any(vars == "CAAobs_ymafr") && length(dat@Dfishery@CAAobs_ymafr)) {
    res$CAAobs_ymafr <- array(NA, dim(dat@Dfishery@CAAobs_ymafr))
    pred <- apply(object@report$CN_ymafrs, 1:5, sum)
    for(y in 1:dat@Dmodel@ny) {
      for(m in 1:dat@Dmodel@nm) {
        for(f in 1:dat@Dfishery@nf) {
          for(r in 1:dat@Dmodel@nr) {
            pred_i <- pred[y, m, , f, r]
            res$CAAobs_ymafr[y, m, , f, r] <- resid_comp(
              obs = dat@Dfishery@CAAobs_ymafr[y, m, , f, r],
              pred = pred_i,
              like = dat@Dfishery@fcomp_like,
              N = dat@Dfishery@CAAN_ymfr[y, m, f, r],
              theta = dat@Dfishery@CAAtheta_f[f],
              stdev = sum(pred_i)/pred_i
            )
          }
        }
      }
    }
  }

  # Catch at length ----
  if (any(vars == "CALobs_ymlfr") && length(dat@Dfishery@CALobs_ymlfr)) {
    res$CALobs_ymlfr <- array(NA, dim(dat@Dfishery@CALobs_ymlfr))
    pred <- apply(object@report$CN_ymlfrs, 1:5, sum)
    for(y in 1:dat@Dmodel@ny) {
      for(m in 1:dat@Dmodel@nm) {
        for(f in 1:dat@Dfishery@nf) {
          for(r in 1:dat@Dmodel@nr) {
            pred_i <- pred[y, m, , f, r]
            res$CALobs_ymlfr[y, m, , f, r] <- resid_comp(
              obs = dat@Dfishery@CALobs_ymlfr[y, m, , f, r],
              pred = pred_i,
              like = dat@Dfishery@fcomp_like,
              N = dat@Dfishery@CALN_ymfr[y, m, f, r],
              theta = dat@Dfishery@CALtheta_f[f],
              stdev = sum(pred_i)/pred_i
            )
          }
        }
      }
    }
  }

  # Indices ----
  if (any(vars == "Iobs_ymi") && length(dat@Dsurvey@Iobs_ymi)) {
    res$Iobs_ymi <- log(dat@Dsurvey@Iobs_ymi/object@report$I_ymi)
    if (type == "pearson") {
      res$Iobs_ymi <- res$Iobs_ymi/dat@Dsurvey@Isd_ymi
    }
  }

  # Indices at age ----
  if (any(vars == "IAAobs_ymai") && length(dat@Dsurvey@IAAobs_ymai)) {
    res$IAAobs_ymai <- array(NA, dim(dat@Dsurvey@IAAobs_ymai))
    pred <- apply(object@report$IN_ymais, 1:4, sum)
    for(y in 1:dat@Dmodel@ny) {
      for(m in 1:dat@Dmodel@nm) {
        for(i in 1:dat@Dsurvey@ni) {
          pred_i <- pred[y, m, , i]
          res$IAAobs_ymai[y, m, , i] <- resid_comp(
            obs = dat@Dsurvey@IAAobs_ymai[y, m, , i],
            pred = pred_i,
            like = dat@Dsurvey@icomp_like,
            N = dat@Dsurvey@IAAN_ymi[y, m, i],
            theta = dat@Dsurvey@IAAtheta_i[i],
            stdev = sum(pred_i)/pred_i
          )
        }
      }
    }
  }

  # Indices at length ----
  if (any(vars == "IALobs_ymli") && length(dat@Dsurvey@IALobs_ymli)) {
    res$IALobs_ymli <- array(NA, dim(dat@Dsurvey@IALobs_ymli))
    pred <- apply(object@report$IN_ymlis, 1:4, sum)
    for(y in 1:dat@Dmodel@ny) {
      for(m in 1:dat@Dmodel@nm) {
        for(i in 1:dat@Dsurvey@ni) {
          pred_i <- pred[y, m, , i]
          res$IALobs_ymli[y, m, , i] <- resid_comp(
            obs = dat@Dsurvey@IALobs_ymli[y, m, , i],
            pred = pred_i,
            like = dat@Dsurvey@icomp_like,
            N = dat@Dsurvey@IALN_ymi[y, m, i],
            theta = dat@Dsurvey@IALtheta_i[i],
            stdev = sum(pred_i)/pred_i
          )
        }
      }
    }
  }

  # Stock composition (at age) ----
  if (any(vars == "SC_ymafrs") && length(dat@Dfishery@SC_ymafrs)) {
    SC_ymafrs <- dat@Dfishery@SC_ymafrs
    res$SC_ymafrs <- array(NA, dim(dat@Dfishery@SC_ymafrs))
    pred <- object@report$SCpred_ymafrs
    for(y in 1:dat@Dmodel@ny) {
      for(m in 1:dat@Dmodel@nm) {
        for(a in 1:dim(SC_ymafrs)[3]) {
          for(f in 1:dim(SC_ymafrs)[4]) {
            for(r in 1:dat@Dmodel@nr) {
              res$SC_ymafrs[y, m, a, f, r, ] <- resid_comp(
                obs = SC_ymafrs[y, m, a, f, r, ],
                pred = pred[y, m, a, f, r, ],
                like = dat@Dfishery@SC_like,
                N = dat@Dfishery@SCN_ymafr[y, m, a, f, r],
                theta = dat@Dfishery@SCtheta_f[f],
                stdev = dat@Dfishery@SCstdev_ymafrs[y, m, a, f, r, ]
              )
            }
          }
        }
      }
    }
  }


  # Stock composition (at age) ----
  if (any(vars == "SC_ymafrs") && length(dat@Dfishery@SC_ymafrs)) {
    SC_ymafrs <- dat@Dfishery@SC_ymafrs
    res$SC_ymafrs <- array(NA, dim(SC_ymafrs))
    pred <- object@report$SCpred_ymafrs
    for(y in 1:dat@Dmodel@ny) {
      for(m in 1:dat@Dmodel@nm) {
        for(a in 1:dim(SC_ymafrs)[3]) {
          for(f in 1:dim(SC_ymafrs)[4]) {
            for(r in 1:dat@Dmodel@nr) {
              res$SC_ymafrs[y, m, a, f, r, ] <- resid_comp(
                obs = SC_ymafrs[y, m, a, f, r, ],
                pred = pred[y, m, a, f, r, ],
                like = dat@Dfishery@SC_like,
                N = dat@Dfishery@SCN_ymafr[y, m, a, f, r],
                theta = dat@Dfishery@SCtheta_f[f],
                stdev = dat@Dfishery@SCstdev_ymafrs[y, m, a, f, r, ]
              )
            }
          }
        }
      }
    }
  }

  # Tag movement ----
  if (any(vars == "tag_ymarrs") && length(dat@Dtag@tag_ymarrs)) {
    tag_ymarrs <- dat@Dtag@tag_ymarrs
    res$tag_ymarrs <- array(NA, dim(tag_ymarrs))
    pred <- object@report$tagpred_ymarrs
    for(y in 1:dim(tag_ymarrs)[1]) {
      for(m in 1:dat@Dmodel@nm) {
        for(a in 1:dim(tag_ymarrs)[3]) {
          for(rf in 1:dat@Dmodel@nr) {
            for(s in 1:dat@Dmodel@ns) {
              res$tag_ymarrs[y, m, a, rf, , s] <- resid_comp(
                obs = tag_ymarrs[y, m, a, rf, , s],
                pred = pred[y, m, a, rf, , s],
                like = dat@Dtag@tag_like,
                N = dat@Dtag@tagN_ymars[y, m, a, rf, s],
                theta = dat@Dtag@tagtheta_s[s],
                stdev = dat@Dtag@tagstdev_s[s]
              )
            }
          }
        }
      }
    }
  }

  return(res)
}
)

# Vectors
resid_comp <- function(obs, pred, like, ...) {
  like <- match.arg(like, choices = eval(formals(like_comp)$type))

  dots <- list(...)

  obs_prob <- obs/sum(obs, na.rm = TRUE)
  pred_prob <- pred/sum(pred, na.rm = TRUE)

  num <- denom <- rep(NA, length(obs))

  # Observed minus predicted
  num[] <- switch(
    like,
    "multinomial" = obs_prob - pred_prob,
    "dirmult1" = obs_prob - pred_prob,
    "dirmult2" = obs_prob - pred_prob,
    "lognormal" = ifelse(obs > 0, log(obs_prob/pred_prob), NA),
    NA
  )

  if (like == "logitnormal") {
    i_fit <- obs > 0
    if (any(i_fit)) {
      i_ref <- rep(FALSE, length(obs))
      i_ref[which(i_fit)[1]] <- TRUE
      num[] <- log(obs_prob/obs_prob[i_ref]) - log(pred_prob/pred_prob[i_ref])
    }
  }

  # Variance
  # Multinomial variance of obs is Np(1-p), variance of pred is p(1-p)/N
  denom[] <- switch(
    like,
    "multinomial" = pred_prob * (1 - pred_prob)/dots$N,
    "dirmult1" = dots$N * pred_prob * (1 - pred_prob) * dots$N * (1 + dots$theta) / (1 + dots$theta * dots$N),
    "dirmult2" = dots$N * pred_prob * (1 - pred_prob) * (dots$N + dots$theta) / (1 + dots$theta),
    "lognormal" = dots$stdev * dots$stdev,
    "logitnormal" = dots$stdev * dots$stdev,
    NA
  )

  res <- num/sqrt(denom)
  return(res)
}

plot_resid_Cobs <- function(fit, f = 1, ...) {
  vars <- "Cobs_ymfr"

  dat <- get_MSAdata(fit)

  res <- residuals(fit, vars = vars, ...)
  x <- apply(res[[vars]][, , f, , drop = FALSE], c(1, 2, 4), identity)

  name <- dat@Dlabel@region

  year <- make_yearseason(dat@Dlabel@year, dat@Dmodel@nm)
  x <- collapse_yearseason(x)

  color <- make_color(ncol(x), type = "region")
  fname <- dat@Dlabel@fleet[f]

  if (!all(is.na(x))) make_tinyplot(year, x, ylab = paste(fname, "catch residual"), name, color, ylim = NULL)
  #matplot(year, x, xlab = "Year", ylab = paste(fname, "catch residual"), type = "o",
  #        col = color, pch = 16, lty = 1)
  #abline(h = 0, lty = 2, col = "grey60")
  #if (ncol(x) > 1) legend("topleft", legend = name, col = color, lwd = 1, pch = 16, horiz = TRUE)

  invisible()
}

plot_resid_Iobs <- function(fit, i = 1, ...) {
  vars <- "Iobs_ymi"

  dat <- get_MSAdata(fit)

  res <- residuals(fit, vars = vars, ...)[[vars]]
  if (is.null(res)) return(invisible())

  x <- apply(res[, , i, drop = FALSE], 1:2, identity)

  name <- dat@Dlabel@index[i]

  year <- make_yearseason(dat@Dlabel@year, dat@Dmodel@nm)
  x <- collapse_yearseason(x)

  if (!all(is.na(x))) {
    plot(year, x, type = "n", xlab = "Year", ylab = paste(name, "residual"))
    lines(year[!is.na(x)], x[!is.na(x)], col = "grey70", lty = 2)
    points(year, x, type = "o", pch = 16)
    abline(h = 0, lty = 1, col = "grey60")
  }


  invisible()
}

plot_resid_CAA <- function(fit, f = 1, r = 1, do_hist = FALSE, ...) {
  vars <- "CAAobs_ymafr"

  dat <- get_MSAdata(fit)

  res <- residuals(fit, vars = vars, ...)[[vars]]
  if (is.null(res)) return(invisible())

  x <- apply(res[, , , f, r, drop = FALSE], 1:3, identity)

  year <- make_yearseason(dat@Dlabel@year, dat@Dmodel@nm)
  x <- collapse_yearseason(x)

  ind <- rowSums(x, na.rm = TRUE) != 0
  yeardiff <- c(diff(year), diff(year)[1])

  .plot_resid_comp(
    year[ind], dat@Dlabel@age, x[ind, , drop = FALSE],
    xlab = "Year", ylab = "Age", do_hist = do_hist,
    xdiff = yeardiff[ind]
  )
}

plot_resid_CAL <- function(fit, f = 1, r = 1, do_hist = FALSE, ...) {
  vars <- "CALobs_ymlfr"

  dat <- get_MSAdata(fit)

  res <- residuals(fit, vars = vars, ...)[[vars]]
  if (is.null(res)) return(invisible())

  x <- apply(res[, , , f, r, drop = FALSE], 1:3, identity)

  year <- make_yearseason(dat@Dlabel@year, dat@Dmodel@nm)
  x <- collapse_yearseason(x)

  ind <- rowSums(x, na.rm = TRUE) != 0
  yeardiff <- c(diff(year), diff(year)[1])

  .plot_resid_comp(
    year[ind], dat@Dmodel@lmid, x[ind, , drop = FALSE],
    xlab = "Year", ylab = "Length", do_hist = do_hist,
    xdiff = yeardiff[ind]
  )
}

plot_resid_IAA <- function(fit, i = 1, do_hist = FALSE, ...) {
  vars <- "IAAobs_ymai"

  dat <- get_MSAdata(fit)

  res <- residuals(fit, vars = vars, ...)[[vars]]
  if (is.null(res)) return(invisible())

  x <- apply(res[, , , i, drop = FALSE], 1:3, identity)

  year <- make_yearseason(dat@Dlabel@year, dat@Dmodel@nm)
  x <- collapse_yearseason(x)

  ind <- rowSums(x, na.rm = TRUE) != 0
  yeardiff <- c(diff(year), diff(year)[1])

  .plot_resid_comp(
    year[ind], dat@Dlabel@age, x[ind, , drop = FALSE],
    xlab = "Year", ylab = "Age", do_hist = do_hist,
    xdiff = yeardiff[ind]
  )
}

plot_resid_IAL <- function(fit, i = 1, do_hist = FALSE, ...) {
  vars <- "IALobs_ymli"

  dat <- get_MSAdata(fit)

  res <- residuals(fit, vars = vars, ...)[[vars]]
  if (is.null(res)) return(invisible())

  x <- apply(res[, , , i, drop = FALSE], 1:3, identity)

  year <- make_yearseason(dat@Dlabel@year, dat@Dmodel@nm)
  x <- collapse_yearseason(x)

  ind <- rowSums(x, na.rm = TRUE) != 0
  yeardiff <- c(diff(year), diff(year)[1])

  .plot_resid_comp(
    year[ind], dat@Dmodel@lmid, x[ind, , drop = FALSE],
    xlab = "Year", ylab = "Length", do_hist = do_hist,
    xdiff = yeardiff[ind]
  )
}

plot_resid_SC <- function(fit, a = 1, f = 1, r = 1, do_hist = FALSE, ...) {
  vars <- "SC_ymafrs"
  dat <- get_MSAdata(fit)

  res <- residuals(fit, vars = vars, ...)[[vars]]
  if (is.null(res)) return(invisible())

  x <- apply(res[, , a, f, r, , drop = FALSE], c(1, 2, 6), identity)

  year <- make_yearseason(dat@Dlabel@year, dat@Dmodel@nm)
  x <- collapse_yearseason(x)

  ind <- rowSums(x, na.rm = TRUE) != 0
  yeardiff <- c(diff(year), diff(year)[1])

  .plot_resid_comp(
    year[ind], dat@Dlabel@stock, x[ind, , drop = FALSE],
    xlab = "Year", ylab = "Stock", do_hist = do_hist,
    xdiff = yeardiff[ind]
  )
}

plot_resid_tagmov <- function(fit, yy = 1, aa = 1, s = 1, ...) {
  vars <- "tag_ymarrs"
  dat <- get_MSAdata(fit)

  res <- residuals(fit, vars = vars, ...)[[vars]]
  if (is.null(res)) return(invisible())

  x <- apply(res[yy, , aa, , , s, drop = FALSE], c(2, 4, 5), identity)

  .plot_resid_comp(z = x, do_hist = TRUE)
}

#' @importFrom graphics hist box
.plot_resid_comp <- function(x = 1:nrow(z), y = 1:ncol(z), z, xlab = "Year", ylab = "Age", zmax = 2,
                             do_hist = FALSE, ydiff, xdiff) {

  if (all(is.na(z))) return(invisible())
  if (do_hist) return(hist(z, xlab = "Residuals", main = ""))

  #if (is.character(y)) {
  #  yval <- y
  #  y <- 1:ncol(z)
  #  yaxt <- "n"
  #} else {
  #  yaxt <- "s"
  #}

  if (is.character(y)) {
    ynum <- 1:length(y)
    tick_y <- function(i) y[i]
    yaxb <- 1:length(y)
  } else {
    ynum <- y
    yaxb <- NULL
    tick_y <- NULL
  }

  #zmax <- min(zmax, max(abs(z), na.rm = TRUE))
  zz <- pmin(z, zmax) |> pmax(-zmax) |> round(2)
  #zlegend <- seq(-zmax, zmax, 0.01) |> round(2)
  #cols <- hcl.colors(length(zlegend), palette = "Blue-Red 2", alpha = 1) |>
  #  structure(names = zlegend)

  if (missing(ydiff)) {
    ydiff <- diff(ynum)
    ydiff <- c(ydiff, ydiff[1])
  }

  if (missing(xdiff)) {
    xdiff <- diff(x)
    xdiff <- c(xdiff, xdiff[1])
  }

  border <- ifelse(any(xdiff < 0.5, na.rm = TRUE), NA, "grey60")
  rect_diff <- ifelse(any(xdiff < 0.5, na.rm = TRUE), 0.475, 0.5)

  df <- lapply(1:nrow(z), function(i) {
    data.frame(
      xleft = x[i] - rect_diff * xdiff[i],
      xright = x[i] + rect_diff * xdiff[i],
      ybottom = ynum - rect_diff * ydiff,
      ytop = ynum + rect_diff * ydiff,
      value = zz[i, ]
      #col = cols[match(zz[i, ], zlegend)]
    )
  })
  df <- do.call(rbind, df)
  df <- df[!is.na(df$value), ]

  #legend_val <- zlegend[c(1, c(0.25, 0.5, 0.75, 1) * length(zlegend))]
  #legend_col <-
  #round(seq(-1, 1, 0.5) * zmax, 2)

  tinyplot(xmin = df$xleft, xmax = df$xright, ymin = df$ybottom, ymax = df$ytop,
           by = df$value, xlab = xlab, ylab = ylab, #yaxs = "i", xaxs = "i",
           bg = "by",
           col = border,
           yaxl = tick_y, yaxb = yaxb,
           #legend = bquote(legend(title = NULL, legend = .(legend_va), 2), pt.bg = c(cols))),
           legend = legend(title = NULL),
           palette = "Blue-Red 2",
           #palette = function(n) hcl.colors(100, palette = "Blue-Red 2", alpha = 0.75, rev = TRUE),
           type = "rect")
  box()

  #plot(NULL, type = "n", xlab = xlab, ylab = ylab, xlim = range(x), ylim = range(y), yaxt = yaxt)
  #sapply(1:nrow(z), function(i) {
  #  if (any(!is.na(zz[i, ]))) {
  #    rect(xleft = x[i] - rect_diff * xdiff[i], xright = x[i] + rect_diff * xdiff[i],
  #         ybottom = y - rect_diff * ydiff, ytop = y + rect_diff * ydiff,
  #         col = cols[match(zz[i, ], zlegend)],
  #         border = border)
  #  }
  #})
  #legend("topleft", legend = c(-zmax, -zmax/2, 0, zmax/2, zmax),
  #       col = "grey60", pt.cex = 1, pt.bg = cols[zlegend %in% c(-zmax, -zmax/2, 0, zmax/2, zmax)],
  #       pch = 22, horiz = TRUE)
  #if (yaxt == "n") {
  #  axis(2, at = y, labels = yval)
  #}
  #box()

  invisible()
}

