
# Data plots ----
#' @name plot-MSA-data
#' @title Plotting functions for data in MSA model
#' @description A set of functions to plot data variables and predicted values (catch, age composition, etc.)
#' @return Invisible data frame of observed and predicted values plotted in base graphics figures
NULL

#' @rdname plot-MSA-data
#' @aliases plot_catch
#'
#' @param fit [MSAassess-class] object returned by [fit_MSA()]
#' @param by Character to indicate dimension for multivariate plots
#' @param f Integer for the corresponding fleet
#' @param prop Logical, whether to plot proportions (TRUE) or absolute numbers
#' @param annual Logical, whether to plot annual values (summed over seasons)
#' @param figure Logical, whether to generate the plot
#' @details
#' - `plot_catch` plots the fishery catch by stock or region (either whole numbers or proportions)
#'
#' @export
plot_catch <- function(fit, f = 1, by = c("region", "stock"), prop = FALSE, annual = FALSE, figure = TRUE) {
  by <- match.arg(by)
  var <- "CB_ymfrs"

  Dlabel <- get_MSAdata(fit)@Dlabel
  year <- Dlabel@year
  nm <- max(length(Dlabel@season), 1)

  if (by == "stock") {
    x <- apply(fit@report[[var]][, , f, , , drop = FALSE], c(1, 2, 5), sum) # C_yms
    name <- Dlabel@stock
  } else {
    x <- apply(fit@report[[var]][, , f, , , drop = FALSE], c(1, 2, 4), sum)  # C_ymr
    name <- Dlabel@region
  }

  if (annual) {
    x <- apply(x, c(1, 3), sum)
  } else {
    year <- make_yearseason(year, nm)
    x <- collapse_yearseason(x)
  }

  fname <- Dlabel@fleet[f]
  if (figure) {
    color <- make_color(ncol(x), type = by)
    ylab <- if (prop) paste(fname, "catch proportion") else paste(fname, "catch")
    barplot2(x, cols = color, leg.names = name, xval = year, ylab = ylab, prop = prop)
  }

  output <- data.frame(
    year = year,
    catch = x,
    fleet = fname
  )

  invisible(output)
}

#' @rdname plot-MSA-data
#' @aliases plot_index
#' @param i Integer, indexes the survey
#' @param zoom Logical, for `plot_index()`. If \code{TRUE}, plots a subset of years with observed data points. Otherwise, plots
#' predicted values over all model years.
#' @param figure Logical, whether to generate the plot
#' @details
#' - `plot_index` plots indices of abundance
#' @export
plot_index <- function(fit, i = 1, zoom = FALSE, figure = TRUE) {
  dat <- get_MSAdata(fit)
  Dlabel <- dat@Dlabel
  year <- Dlabel@year
  ny <- dat@Dmodel@ny
  nm <- dat@Dmodel@nm

  iname <- Dlabel@index[i]

  iobs <- dat@Dsurvey@Iobs_ymi[, , i, drop = FALSE]
  isd <- dat@Dsurvey@Isd_ymi[, , i, drop = FALSE]

  ni_plot <- length(i)

  ipred <- fit@report$I_ymi[, , i, drop = FALSE]

  year <- make_yearseason(year, nm)
  ipred <- collapse_yearseason(ipred)
  iobs <- collapse_yearseason(iobs)
  isd <- collapse_yearseason(isd)

  ind_all <- apply(iobs, 1, function(x) !all(is.na(x)))

  output <- NULL
  if (any(!is.na(iobs))) {
    iupper <- exp(log(iobs) + 1.96 * isd)
    ilower <- exp(log(iobs) - 1.96 * isd)

    output <- lapply(1:ni_plot, function(ii) {
      .output <- data.frame(
        year = year, obs = iobs[, ii], pred = ipred[, ii], lwr = ilower[, ii], upr = iupper[, ii], name = iname[ii]
      )
      if (zoom) {
        mind <- rep(1:nm, ny)
        mobs <- unique(mind[!is.na(.output$obs)])
        ind <- mind %in% mobs
        return(.output[ind, ])
      } else {
        return(.output)
      }
    })
    output <- do.call(rbind, output)

    if (figure) {
      irange <- lapply(1:ni_plot, function(ii) {
        data.frame(year = range(year[ind_all]), y = c(0, 1.1) * max(iupper[, ii], ipred[, ii], na.rm = TRUE), name = iname[ii])
      })
      irange <- do.call(rbind, irange)

      irange$name <- factor(irange$name, iname)
      output$name <- factor(output$name, iname)

      tinyplot(
        irange$year, irange$y,
        facet = irange$name, facet.args = list(free = TRUE),
        type = "n", grid = TRUE,
        xlab = "Year",  ylab = "Index"
      )
      tinyplot(
        output$year, output$pred,
        facet = output$name, facet.args = list(free = TRUE),
        type = "l", col = "red",
        add = TRUE
      )
      tinyplot(
        output$year, output$obs,
        facet = output$name, facet.args = list(free = TRUE),
        ymin = output$lwr, ymax = output$upr,
        type = "pointrange", pch = 16,
        add = TRUE
      )

      #plot(year, iobs, xlab = "Year", ylab = iname, type = "p", pch = 16,
      #     ylim = c(0, 1.1) * range(ipred, iupper, na.rm = TRUE), zero_line = TRUE)
      #arrows(year, y0 = ilower, y1 = iupper, length = 0)
      #lines(year, ipred, lwd = 2, col = 2, type = ifelse(length(year) > 10, "l", "o"))
    }
  }

  invisible(output)
}


#' @rdname plot-MSA-data
#' @aliases plot_CAA
#' @param f Integer, indexes the fleet
#' @param r Integer, indexes the region
#' @param do_mean Logical, whether to plot full compositions or time series of mean length or mean age
#' @details
#' - `plot_CAA` plots the fishery catch at age
#' @importFrom stats weighted.mean
#' @export
plot_CAA <- function(fit, f = 1, r = 1, do_mean = FALSE, figure = TRUE) {
  dat <- get_MSAdata(fit)
  output <- NULL

  if (sum(dat@Dfishery@CAAN_ymfr, na.rm = TRUE)) {
    N <- apply(dat@Dfishery@CAAN_ymfr[, , f, r, drop = FALSE], 1:2, identity) |> t() |> as.numeric()
    N[is.na(N)] <- 0

    if (any(N > 0)) {
      Dlabel <- dat@Dlabel
      nm <- max(length(Dlabel@season), 1)

      year <- Dlabel@year
      year <- make_yearseason(year, nm)

      pred <- apply(fit@report$CN_ymafrs[, , , f, r, , drop = FALSE], 1:3, sum)
      pred <- collapse_yearseason(pred) |> apply(1, function(x) x/sum(x, na.rm = TRUE)) |> t()
      pred[is.na(pred)] <- 0

      obs <- apply(dat@Dfishery@CAAobs_ymafr[, , , f, r, drop = FALSE], 1:3, identity)
      obs <- collapse_yearseason(obs) |> apply(1, function(x) x/sum(x, na.rm = TRUE)) |> t()

      include <- rowSums(obs, na.rm = TRUE) > 0
      if (do_mean) {
        age <- seq(1, dat@Dmodel@na) - 1
        mpred <- apply(pred, 1, function(w) weighted.mean(x = age, w = w))
        mobs <- apply(obs, 1, function(w) weighted.mean(x = age, w = w))

        if (figure) {
          ylim <- c(0.9, 1.1) * range(c(mobs, mpred), na.rm = TRUE)

          plot(year[include], mobs[include], xlab = "Year", ylab = "Mean age", pch = 1, type = "o", ylim = ylim)
          lines(year[include], mpred[include], col = 2, lwd = 2)
        }

        output <- data.frame(
          year = year,
          pred = mpred,
          obs = mobs,
          fleet = dat@Dlabel@fleet[f],
          region = dat@Dlabel@region[r]
        )
      } else {
        if (figure) {
          plot_composition(obs[include, , drop = FALSE], pred[include, , drop = FALSE], xval = Dlabel@age, ylab = "Proportion",
                           zval = year[include], N = N[include])
        }

        obs_df <- structure(obs, dimnames = list(year = year, age = Dlabel@age)) |>
          reshape2::melt(value.name = "obs")
        pred_df <- structure(pred, dimnames = list(year = year, age = Dlabel@age)) |>
          reshape2::melt(value.name = "pred")
        N_df <- data.frame(
          year = year,
          fleet = dat@Dlabel@fleet[f],
          region = dat@Dlabel@region[r],
          N = N
        )
        output <- merge(N_df, merge(obs_df, pred_df))
      }
    }
  }

  invisible(output)
}

#' @rdname plot-MSA-data
#' @aliases plot_CAL
#' @param agg Logical, whether to aggregate composition across all time steps
#' @details
#' - `plot_CAL` plots the catch at length
#' @importFrom stats aggregate
#' @export
plot_CAL <- function(fit, f, r, agg = FALSE, do_mean = FALSE, figure = TRUE) {
  dat <- get_MSAdata(fit)
  output <- NULL
  if (missing(f)) f <- 1:dat@Dfishery@nf
  if (missing(r)) r <- 1:dat@Dmodel@nr
  nf_plot <- length(f)
  nr_plot <- length(r)

  if (sum(dat@Dfishery@CALN_ymfr, na.rm = TRUE)) {
    N <- dat@Dfishery@CALN_ymfr[, , f, r, drop = FALSE]
    N[is.na(N)] <- 0

    if (any(N > 0)) {
      Dlabel <- dat@Dlabel
      nm <- max(length(Dlabel@season), 1)

      year <- Dlabel@year
      year <- make_yearseason(year, nm)

      pred <- apply(fit@report$CN_ymlfrs[, , , f, r, , drop = FALSE], 1:5, sum)
      pred <- collapse_yearseason(pred) |> apply(c(1, 3, 4), function(x) x/sum(x, na.rm = TRUE)) |> aperm(c(2, 1, 3, 4))
      pred[is.na(pred)] <- 0

      obs <- dat@Dfishery@CALobs_ymlfr[, , , f, r, drop = FALSE]
      obs <- collapse_yearseason(obs) |> apply(c(1, 3, 4), function(x) x/sum(x, na.rm = TRUE)) |> aperm(c(2, 1, 3, 4))

      N <- collapse_yearseason(N)

      if (do_mean) {
        mpred <- apply(pred, c(1, 3, 4), function(w) weighted.mean(x = dat@Dmodel@lmid, w = w))
        mobs <- apply(obs, c(1, 3, 4), function(w) weighted.mean(x = dat@Dmodel@lmid, w = w))
        Nobs <- apply(mobs, 2:3, sum, na.rm = TRUE)

        output <- lapply(1:nf_plot, function(ff) {
          .m <- lapply(1:nr_plot, function(rr) {
            .output <- output[output$fleet == dat@Dlabel@fleet[ff] & output$region == dat@Dlabel@region[rr], ]
            if (Nobs[ff, rr]) {
              data.frame(
                year = year,
                obs = mobs[, ff, rr],
                pred = mpred[, ff, rr],
                fleet = dat@Dlabel@fleet[ff],
                region = dat@Dlabel@region[rr]
              )
            } else {
              data.frame()
            }
          })
          do.call(rbind, .m)
        })
        output <- do.call(rbind, output)

        if (figure) {

          mrange <- lapply(1:nf_plot, function(ff) {
            .m <- lapply(1:nr_plot, function(rr) {
              if (Nobs[ff, rr]) {
                .output <- output[output$fleet == dat@Dlabel@fleet[ff] & output$region == dat@Dlabel@region[rr], ]
                data.frame(
                  year = range(year),
                  y = c(0.9, 1.1) * range(.output$obs, .output$pred, na.rm = TRUE),
                  fleet = dat@Dlabel@fleet[ff],
                  region = dat@Dlabel@region[rr]
                )
              } else {
                data.frame()
              }
            })
            do.call(rbind, .m)
          })
          mrange <- do.call(rbind, mrange)

          tinyplot(
            mrange$year, mrange$y,
            facet = paste(mrange$fleet, mrange$region, sep = ", "), facet.args = list(free = TRUE),
            type = "n", grid = TRUE,
            xlab = "Year", ylab = "Mean length"
          )

          out_dat <- output[!is.na(output$obs), ]
          tinyplot(
            out_dat$year, out_dat$obs,
            facet = paste(out_dat$fleet, out_dat$region, sep = ", "),
            facet.args = list(free = TRUE),
            type = "o",
            add = TRUE
          )

          out_pred <- output[!is.na(output$pred), ]
          tinyplot(
            out_pred$year, out_pred$pred,
            facet = paste(out_pred$fleet, out_pred$region, sep = ", "), facet.args = list(free = TRUE),
            type = "l", col = "red", lwd = 2, add = TRUE
          )
        }


      } else {
        obs_df <- structure(
          obs,
          dimnames = list(year = year, lmid = dat@Dmodel@lmid,
                          fleet = dat@Dlabel@fleet[f], region = dat@Dlabel@region[r])
        ) |>
          reshape2::melt(value.name = "obs")
        pred_df <- structure(
          pred,
          dimnames = list(year = year, lmid = dat@Dmodel@lmid,
                          fleet = dat@Dlabel@fleet[f], region = dat@Dlabel@region[r])
        ) |>
          reshape2::melt(value.name = "pred")
        N_df <- structure(
          N,
          dimnames = list(year = year, fleet = dat@Dlabel@fleet[f], region = dat@Dlabel@region[r])
        ) |>
          reshape2::melt(value.name = "N")

        output <- merge(N_df, merge(obs_df, pred_df))

        if (agg) {

          .obs_agg <- aggregate(
            output$N * output$obs, by = list(lmid = output$lmid, fleet = output$fleet),
            FUN = sum, na.rm = TRUE
          )
          names(.obs_agg)[3] <- "obs"
          .pred_agg <- aggregate(
            output$N * output$pred, by = list(lmid = output$lmid, fleet = output$fleet),
            FUN = sum, na.rm = TRUE
          )
          names(.pred_agg)[3] <- "pred"

          output <- merge(.obs_agg, .pred_agg, sort = FALSE)

          if (figure) {
            frange <- lapply(unique(output$fleet), function(ff) {
              .output <- output[output$fleet == ff, ]
              data.frame(
                lmid = range(output$lmid),
                y = c(0, 1.1) * max(.output$obs, .output$pred, na.rm = TRUE),
                fleet = ff
              )
            })
            frange <- do.call(rbind, frange)
            tinyplot(
              frange$lmid, frange$y,
              facet = frange$fleet, facet.args = list(free = TRUE),
              type = "n",
              xlab = "Length", ylab = "N"
            )
            tinyplot(
              output$lmid, output$obs,
              facet = output$fleet, facet.args = list(free = TRUE),
              type = "area", pch = 16,
              add = TRUE
            )
            tinyplot(
              output$lmid, output$pred,
              facet = output$fleet, facet.args = list(free = TRUE),
              type = "l", col = 2, lwd = 3,
              add = TRUE
            )
          }

        } else if (figure) {

          if (length(f) > 1 || length(r) > 1) stop("Specify one value each for f and r")

          include <- apply(obs, 1, sum, na.rm = TRUE) > 0
          plot_composition(
            matrix(obs[include, , , ], sum(include), dat@Dmodel@nl),
            matrix(pred[include, , , ], sum(include), dat@Dmodel@nl),
            xval = dat@Dmodel@lmid,
            xlab = "Length", ylab = "Proportion",
            zval = year[include], N = as.numeric(N[include, , ])
          )
        }
      }
    }
  }

  invisible(output)
}

#' @rdname plot-MSA-data
#' @aliases plot_IAA
#' @details
#' - `plot_IAA` plots the index age composition
#' @export
plot_IAA <- function(fit, i = 1, do_mean = FALSE, figure = TRUE) {
  dat <- get_MSAdata(fit)
  output <- NULL

  if (sum(dat@Dsurvey@IAAN_ymi, na.rm = TRUE)) {
    N <- apply(dat@Dsurvey@IAAN_ymi[, , i, drop = FALSE], 1:2, identity) |> t() |> as.numeric()
    N[is.na(N)] <- 0

    if (any(N > 0)) {
      Dlabel <- dat@Dlabel
      nm <- max(length(Dlabel@season), 1)

      year <- Dlabel@year
      year <- make_yearseason(year, nm)

      pred <- apply(fit@report$IN_ymais[, , , i, , drop = FALSE], 1:3, sum)
      pred <- collapse_yearseason(pred) |> apply(1, function(x) x/sum(x, na.rm = TRUE)) |> t()
      pred[is.na(pred)] <- 0

      obs <- apply(dat@Dsurvey@IAAobs_ymai[, , , i, drop = FALSE], 1:3, identity)
      obs <- collapse_yearseason(obs) |> apply(1, function(x) x/sum(x, na.rm = TRUE)) |> t()

      include <- rowSums(obs, na.rm = TRUE) > 0
      if (do_mean) {
        age <- seq(1, dat@Dmodel@na) - 1
        mpred <- apply(pred, 1, function(w) weighted.mean(x = age, w = w))
        mobs <- apply(obs, 1, function(w) weighted.mean(x = age, w = w))

        if (figure) {
          ylim <- c(0.9, 1.1) * range(c(mobs, mpred), na.rm = TRUE)

          plot(year[include], mobs[include], xlab = "Year", ylab = "Mean age", pch = 1, type = "o", ylim = ylim)
          lines(year[include], mpred[include], col = 2, lwd = 2)
        }

        output <- data.frame(
          year = year,
          pred = mpred,
          obs = mobs,
          name = dat@Dlabel@index[i]
        )

      } else {
        if (figure) {
          plot_composition(obs[include, , drop = FALSE], pred[include, , drop = FALSE], xval = Dlabel@age, ylab = "Proportion",
                           zval = year[include], N = N[include])
        }

        obs_df <- structure(obs, dimnames = list(year = year, age = Dlabel@age)) |>
          reshape2::melt(value.name = "obs")
        pred_df <- structure(pred, dimnames = list(year = year, age = Dlabel@age)) |>
          reshape2::melt(value.name = "pred")
        N_df <- data.frame(
          year = year,
          name = dat@Dlabel@index[i],
          N = N
        )
        output <- merge(N_df, merge(obs_df, pred_df))
      }

    }
  }

  invisible(output)
}

#' @rdname plot-MSA-data
#' @aliases plot_IAL
#' @details
#' - `plot_IAL` plots the index length composition
#' @export
plot_IAL <- function(fit, i = 1, do_mean = FALSE, figure = TRUE) {
  dat <- get_MSAdata(fit)
  output <- NULL

  if (sum(dat@Dsurvey@IALN_ymi, na.rm = TRUE)) {
    N <- apply(dat@Dsurvey@IALN_ymi[, , i, drop = FALSE], 1:2, identity) |> t() |> as.numeric()
    N[is.na(N)] <- 0

    if (any(N > 0)) {
      Dlabel <- dat@Dlabel
      nm <- max(length(Dlabel@season), 1)

      year <- Dlabel@year
      year <- make_yearseason(year, nm)

      pred <- apply(fit@report$IN_ymlis[, , , i, , drop = FALSE], 1:3, sum)
      pred <- collapse_yearseason(pred) |> apply(1, function(x) x/sum(x, na.rm = TRUE)) |> t()
      pred[is.na(pred)] <- 0

      obs <- apply(dat@Dsurvey@IALobs_ymli[, , , i, drop = FALSE], 1:3, identity)
      obs <- collapse_yearseason(obs) |> apply(1, function(x) x/sum(x, na.rm = TRUE)) |> t()

      include <- rowSums(obs, na.rm = TRUE) > 0
      if (do_mean) {
        mpred <- apply(pred, 1, function(w) weighted.mean(x = dat@Dmodel@lmid, w = w))
        mobs <- apply(obs, 1, function(w) weighted.mean(x = dat@Dmodel@lmid, w = w))

        if (figure) {
          ylim <- c(0.9, 1.1) * range(c(mobs, mpred), na.rm = TRUE)

          plot(year[include], mobs[include], xlab = "Year", ylab = "Mean length", pch = 1, type = "o", ylim = ylim)
          lines(year[include], mpred[include], col = 2, lwd = 2)
        }

        output <- data.frame(
          year = year,
          pred = mpred,
          obs = mobs,
          name = dat@Dlabel@index[i]
        )

      } else {
        if (figure) {
          plot_composition(obs[include, ], pred[include, ], xval = dat@Dmodel@lmid,
                           xlab = "Length", ylab = "Proportion",
                           zval = year[include > 0], N = N[include])
        }

        obs_df <- structure(obs, dimnames = list(year = year, lmid = dat@Dmodel@lmid)) |>
          reshape2::melt(value.name = "obs")
        pred_df <- structure(pred, dimnames = list(year = year, lmid = dat@Dmodel@lmid)) |>
          reshape2::melt(value.name = "pred")
        N_df <- data.frame(
          year = year,
          name = dat@Dlabel@index[i],
          N = N
        )
        output <- merge(N_df, merge(obs_df, pred_df))
      }

    }
  }

  invisible()
}

#' @rdname plot-MSA-data
#' @param ff Integer, indexes the aggregate fleet (for stock composition data)
#' @param aa Integer, indexes the aggregate age class (for stock composition and tag data)
#' @aliases plot_SC
#' @details
#' - `plot_SC` plots the stock composition
#' @importFrom graphics matlines
#' @export
plot_SC <- function(fit, ff = 1, aa = 1, r = 1, prop = FALSE, figure = TRUE) {
  dat <- get_MSAdata(fit)
  output <- NULL

  if (dat@Dmodel@ns == 1) stop("Stock composition figure not needed.")

  if (sum(dat@Dfishery@SC_ymafrs, na.rm = TRUE)) {
    N <- apply(dat@Dfishery@SC_ymafrs[, , aa, ff, r, , drop = FALSE], 1:2, sum)
    N[is.na(N)] <- 0

    if (any(N > 0)) {
      Dlabel <- dat@Dlabel
      nm <- max(length(Dlabel@season), 1)

      year <- Dlabel@year
      year <- make_yearseason(year, nm)

      pred <- apply(fit@report$SCpred_ymafrs[, , aa, ff, r, , drop = FALSE], c(1, 2, 6), identity)
      pred <- collapse_yearseason(pred) |> apply(1, function(x) x/sum(x, na.rm = TRUE)) |> t()
      pred[is.na(pred)] <- 0

      N <- collapse_yearseason(N)

      obs <- apply(dat@Dfishery@SC_ymafrs[, , aa, ff, r, , drop = FALSE], c(1, 2, 6), identity)
      obs <- collapse_yearseason(obs) |> apply(1, function(x) x/sum(x, na.rm = TRUE)) |> t()

      if (figure) {
        if (prop) {
          color <- make_color(ncol(pred), type = "stock")
          barplot2(pred, cols = color, leg.names = Dlabel@stock, xval = year)

          if (dat@Dmodel@ns == 2) {
            obs_cumsum <- apply(obs, 1, cumsum) |> t()
            matlines(obs_cumsum[, -dat@Dmodel@ns, drop = FALSE], type = "o", col = 2, pch = 16, ylim = c(0, 1))
          }
        } else {
          include <- rowSums(obs, na.rm = TRUE) > 0

          plot_composition(obs[include, , drop = FALSE], pred[include, , drop = FALSE], xval = 1:dat@Dmodel@ns,
                           xlab = "Stock", ylab = "Proportion",
                           zval = year[include], N = N[include],
                           xaxislab = Dlabel@stock)
        }
      }

      obs_df <- structure(obs, dimnames = list(year = year, stock = Dlabel@stock)) |>
        reshape2::melt(value.name = "obs")
      pred_df <- structure(pred, dimnames = list(year = year, stock = Dlabel@stock)) |>
        reshape2::melt(value.name = "pred")

      N_df <- data.frame(
        year = year,
        ff = ff,
        aa = aa,
        region = dat@Dlabel@region[r]
      )
      output <- merge(N_df, merge(obs_df, pred_df))

      if (length(dat@Dfishery@SCstdev_ymafrs)) {
        se_df <- apply(dat@Dfishery@SCstdev_ymafrs[, , aa, ff, r, , drop = FALSE], c(1, 2, 6), identity) |>
          collapse_yearseason() |>
          structure(dimnames = list(year = year, stock = Dlabel@stock)) |>
          reshape2::melt(value.name = "se")
        output <- merge(output, se_df)
      }
    }
  }

  invisible(output)
}

#' @importFrom graphics lines mtext axis
plot_composition <- function(obs, pred = NULL, xval = 1:ncol(obs), xlab = "Age",
                             ylab = "Value", zval = 1:nrow(obs), N = rowSums(obs),
                             xaxislab = xval, ncol = 4, nrow = 4) {

  old_par <- par(no.readonly = TRUE)
  on.exit(par(old_par))
  par(mfcol = c(nrow, ncol), mar = rep(0, 4), oma = c(5.1, 5.1, 2.1, 2.1))

  ylim <- c(0, 1.1) * range(pred, obs, na.rm = TRUE)

  yaxp <- c(0, max(pretty(ylim, n = 4)), 4)
  if (max(obs, pred, na.rm = TRUE) == 1) yaxp <- c(0, 1, 4)
  las <- 1
  nplot <- nrow(pred)
  nset <- nrow * ncol
  xaxt_manual <- any(xaxislab != xval)
  N <- round(N, 2)

  for(i in 1:nplot) {
    yaxt <- ifelse(i %% nset %in% c(1:nrow), "s", "n") # TRUE = first column
    xaxt <- ifelse(i < nplot && i %% nrow %in% c(1:(nrow-1)), "n", "s") # TRUE = first three rows

    plot(xval, obs[i, ], type = "o", ylim = ylim, yaxp = yaxp, xaxt = ifelse(xaxt_manual, "n", xaxt), yaxt = yaxt, las = las)
    if (xaxt_manual && xaxt == "s") axis(1, at = xval, labels = xaxislab)
    if (!is.null(pred)) lines(xval, pred[i, ], lwd = 2, col = 2)
    legend("topright", legend = c(zval[i], ifelse(is.null(N), "", paste0("N = ", N[i]))), bty = "n", xjust = 1)

    if (i %% nset == 1) {
      mtext(xlab, side = 1, line = 3, outer = TRUE)
      mtext(ylab, side = 2, line = 3.5, outer = TRUE)
    }
  }

  invisible()
}

#' @rdname plot-MSA-data
#' @param yy Integer, indexes the aggregate years (for the tag data)
#' @param s Integer, indexes the stock
#' @aliases plot_tagmov
#' @details
#' - `plot_tagmov` plots the tag movements
#' @export
plot_tagmov <- function(fit, s = 1, yy = 1, aa = 1, figure = TRUE) {
  dat <- get_MSAdata(fit)
  output <- NULL

  if (dat@Dmodel@nr == 1) stop("Tag movement figure not needed.")

  if (sum(dat@Dtag@tag_ymarrs, na.rm = TRUE)) {
    N <- apply(dat@Dtag@tag_ymarrs[yy, , aa, , , s, drop = FALSE], c(2, 4), sum)
    N[is.na(N)] <- 0

    if (any(N > 0)) {
      Dlabel <- dat@Dlabel

      pred <- apply(fit@report$tagpred_ymarrs[yy, , aa, , , s, drop = FALSE], c(2, 4, 5), identity)
      pred2 <- collapse_yearseason(pred)

      obs <- apply(dat@Dtag@tag_ymarrs[yy, , aa, , , s, drop = FALSE], c(2, 4, 5), identity)
      obs2 <- apply(obs, c(1, 2), function(x) x/sum(x, na.rm = TRUE)) |> aperm(c(2, 3, 1)) |>
        collapse_yearseason() # Season + region of origin
      obs2[is.na(obs2)] <- 0

      N <- collapse_yearseason(N)

      if (figure) {
        nrow <- dat@Dmodel@nr
        ncol <- dat@Dmodel@nm

        z <- paste(
          paste("Origin:", rep(Dlabel@region, dat@Dmodel@nm)),
          paste(rep(paste("\nSeason", 1:dat@Dmodel@nm), each = dat@Dmodel@nr))
        )

        plot_composition(obs2, pred2, xval = 1:dat@Dmodel@nr,
                         xlab = "Destination", ylab = "Proportion",
                         zval = z, N = N,
                         xaxislab = Dlabel@region, ncol = ncol, nrow = nrow)
      }

      obs_df <- structure(
        obs,
        dimnames = list(season = Dlabel@season, from = Dlabel@region, to = Dlabel@region)
      ) |>
        reshape2::melt(value.name = "obs")
      pred_df <- structure(
        pred,
        dimnames = list(season = Dlabel@season, from = Dlabel@region, to = Dlabel@region)
      ) |>
        reshape2::melt(value.name = "pred")

      N_df <- data.frame(
        year = yy,
        yy = yy,
        aa = aa,
        stock = dat@Dlabel@stock[s]
      )
      output <- merge(N_df, merge(obs_df, pred_df))
    }
  }

  invisible(output)
}
