## Author: Ugolkov Yaroslav
## First created: 2025-09-05
## Description: sg-gof-par-dist and its helper functions
## Keywords: SimuRg, sg-gof-par-dist, goodness-of-fit
#'
#' Plot distributions of individual parameters with optional theoretical density
#'@description
#' This function creates histograms of individual parameter estimates,
#' with optional overlay of the theoretical distribution based on the population parameters and OMEGA matrix.
#'
#' @inheritParams sg_dummy
#'
#' @param plot_type Character string specifying the type of plot to generate.
#'   Options: 'DIST' for parameter distributions (default), 'QQ' for Q-Q plots,
#'   'correlations' for correlation matrix plot.
#' @return A `ggplot` object containing histogram(s) of individual parameter distributions, optionally overlaid with theoretical densities.
#'
#' @details
#' The function visualizes the distribution of post hoc individual parameter estimates (from `$PATAB`) and optionally compares them with the expected population-level variability using `$SUMTAB$DISTRIBUTION` (see Details).
#' Shrinkage values are calculated based on `ETAshrinkage_var` from `$SUMTAB` and included in facet labels.
#'
#' Theoretical distributions are generated by sampling from a multivariate normal distribution (mean zero, covariance from `$OMEGAMAT`) and mapping to the individual-parameter scale using typical values  and the `DISTRIBUTION` column: `logNormal` uses `TV * exp(ETA)`, `normal` uses `TV + ETA`, `logitNormal` uses the inverse logit of `logit(TV) + ETA`. Parameters with missing or empty `DISTRIBUTION` do not get a theoretical curve.
#'
#'
#' @examples
#' \donttest{
#' # Basic usage: distribution plots for all parameters
#' fpath_i <- system.file("extdata", "simurg_object", "Warfarin_PK.RData", package = "SimuRg")
#' sg_gof_par_dist(fpath_i = fpath_i)
#'
#' # Distribution plots for selected parameters (natural scale)
#' sg_gof_par_dist(fpath_i = fpath_i, par_seq = c("ka", "Cl"))
#'
#' # Distribution plots with theoretical densities disabled
#' sg_gof_par_dist(fpath_i = fpath_i, tdist = FALSE)
#'
#' # Distribution plots for ETA
#' sg_gof_par_dist(fpath_i = fpath_i, par_seq = c("eta_ka", "eta_Cl"), par_type = "RE")
#'
#' # Q-Q plots for all ETA parameters
#' sg_gof_par_dist(fpath_i = fpath_i, plot_type = "QQ")
#'
#' # Q-Q plot for a specific ETA parameter
#' sg_gof_par_dist(fpath_i = fpath_i, plot_type = "QQ", par_seq = "eta_ka")
#'
#' # Correlation matrix for selected parameters
#' sg_gof_par_dist(fpath_i = fpath_i, plot_type = "correlations")
#' }
#' @import dplyr
#' @importFrom MASS mvrnorm
#' @import ggplot2
#' @importFrom scales pretty_breaks
#' @importFrom grDevices recordPlot
#' @importFrom graphics abline box grid hist lines mtext par points text
#' @importFrom stats lm
#' @export

sg_gof_par_dist <- function(fpath_i, par_seq = NULL,
                            par_type = "Ind", n_bins = 30, tdist = TRUE, plot_type = 'DIST'){
  load_smrg_obj <- function(x) {
    if (inherits(x, "list")) return(x)
    if (!inherits(x, "character") || length(x) != 1L) {
      stop("fpath_i object should be either an sg_fit object (list), or a path to saved sg_fit object")
    }
    if (!file.exists(x)) stop("File does not exist: ", x)

    ext <- tolower(tools::file_ext(x))
    if (ext %in% c("rdata", "rda", "r")) {
      nm <- load(x)
      if (length(nm) == 0) stop("No objects found in file: ", x)
      return(get(nm[1]))
    }

    if (ext == "json") {
      obj <- jsonlite::read_json(x, simplifyVector = TRUE)
      if (is.list(obj) && all(c("GFO", "GCO") %in% names(obj))) obj <- obj$GFO
      return(obj)
    }

    stop("Unsupported file extension: .", ext, " (expected .RData/.rda or .json)")
  }

  smrg_obj <- load_smrg_obj(fpath_i)

  # Coerce core components after JSON import if needed
  if ("OMEGAMAT" %in% names(smrg_obj) && !is.null(smrg_obj$OMEGAMAT) && !is.matrix(smrg_obj$OMEGAMAT)) {
    smrg_obj$OMEGAMAT <- as.matrix(smrg_obj$OMEGAMAT)
  }
  if ("SIGMAMAT" %in% names(smrg_obj) && !is.null(smrg_obj$SIGMAMAT) && !is.matrix(smrg_obj$SIGMAMAT)) {
    smrg_obj$SIGMAMAT <- as.matrix(smrg_obj$SIGMAMAT)
  }
  if ("COVMAT" %in% names(smrg_obj) && !is.null(smrg_obj$COVMAT) && !is.matrix(smrg_obj$COVMAT)) {
    smrg_obj$COVMAT <- as.matrix(smrg_obj$COVMAT)
  }
  if ("CORRMAT" %in% names(smrg_obj) && !is.null(smrg_obj$CORRMAT) && !is.matrix(smrg_obj$CORRMAT)) {
    smrg_obj$CORRMAT <- as.matrix(smrg_obj$CORRMAT)
  }
  if ("PATAB" %in% names(smrg_obj) && !is.null(smrg_obj$PATAB) && !inherits(smrg_obj$PATAB, "data.frame")) {
    smrg_obj$PATAB <- as.data.frame(smrg_obj$PATAB)
  }
  if ("SUMTAB" %in% names(smrg_obj) && !is.null(smrg_obj$SUMTAB) && !inherits(smrg_obj$SUMTAB, "data.frame")) {
    smrg_obj$SUMTAB <- as.data.frame(smrg_obj$SUMTAB)
  }

  prm_i <- smrg_obj$COVMAT
  patab_i <- smrg_obj$PATAB
  omegamat_i <- smrg_obj$OMEGAMAT
  ### Distribution of individual parameters
  if(is.null(par_seq)){
    par_seq <- names(patab_i %>% select(-contains("eta"),-ID))
    # If OMEGAMAT is empty/missing, keep all parameters from PATAB
    if (!is.null(omegamat_i) && is.matrix(omegamat_i) && !is.null(colnames(omegamat_i)) && length(colnames(omegamat_i)) > 0) {
      par_seq <- par_seq[par_seq %in% sub("^omega_", "", colnames(omegamat_i))]
    }
  }

  # --- Validate par_seq ---
  missing_par <- setdiff(par_seq, names(patab_i))
  if (length(missing_par) > 0) {
    warning(paste("Parameters not found in PATAB:", paste(missing_par, collapse = ", ")))
  }

  par_seq <- intersect(par_seq, names(patab_i))
  if (length(par_seq) == 0) {
    return(invisible(NULL))
  }


  if (plot_type == 'correlations') {
    # Prepare data
    if(is.null(par_seq)){
      selected_cols <- setdiff(names(patab_i), "ID")
    } else {
      selected_cols <- par_seq
    }

    ds_eta <- patab_i %>% dplyr::select(ID, dplyr::all_of(selected_cols))
    plot_cols <- setdiff(names(ds_eta), c("ID", "eta_Vd", "eta_CL"))
    n <- length(plot_cols)

    # Save current graphics parameters
    old_par <- par(no.readonly = TRUE)
    on.exit(par(old_par))

    # Create layout for plots
    par(mfrow = c(n, n),
        mar = c(1.5, 1.5, 0.5, 0.5),
        oma = c(2, 2, 2, 1))

    for(i in 1:n) {
      for(j in 1:n) {
        if(i == j) {
          # DIAGONAL - histogram with density
          # First draw empty plot with histogram stats to set proper ylim
          h <- hist(ds_eta[[plot_cols[i]]], plot = FALSE, breaks = 10)

          plot(1, type = "n",
               xlim = range(ds_eta[[plot_cols[i]]]),
               ylim = c(0, max(h$density) * 1.1),
               xlab = "", ylab = "",
               axes = TRUE, cex.axis = 1, mgp = c(2, 0.5, 0))

          # Add grid
          grid(col = "grey80", lty = 1, lwd = 0.5)

          # Add histogram
          hist(ds_eta[[plot_cols[i]]],
               main = "",
               xlab = "",
               ylab = "",
               col = MSDcol[2],
               border = "grey45",
               probability = TRUE,
               breaks = 10,
               add = TRUE)

          box()

          # Density curve
          lines(density(ds_eta[[plot_cols[i]]]),
                col = "grey25",
                lwd = 1.7)

          # Vertical line at 0
          abline(v = 0, col = "grey25", lty = 3, lwd = 1)

        } else if(i < j) {
          # UPPER TRIANGLE - Pearson correlation with significance stars
          plot(1, type = "n",
               xlim = c(0, 1), ylim = c(0, 1),
               xlab = "", ylab = "",
               xaxt = "n", yaxt = "n",
               axes = FALSE)
          box()

          # Pearson correlation
          cor_test <- cor.test(ds_eta[[plot_cols[i]]], ds_eta[[plot_cols[j]]],
                               method = "pearson", use = "complete.obs")
          cor_val <- cor_test$estimate
          p_val <- cor_test$p.value

          # Significance stars
          stars <- ifelse(p_val < 0.001, "***",
                          ifelse(p_val < 0.01, "**",
                                 ifelse(p_val < 0.05, "*", "")))

          # Correlation text with stars
          text(0.5, 0.6,
               labels = paste0("Pearson correlation = ", round(cor_val, 2), stars),
               cex = 1.3,
               font = 2,
               col = "black")

        } else {
          # LOWER TRIANGLE - scatter plot with regression line
          plot(ds_eta[[plot_cols[j]]], ds_eta[[plot_cols[i]]],
               pch = 16,
               col = MSDcol[1],
               xlab = "",
               ylab = "",
               cex = 1,
               axes = TRUE,
               cex.axis = 1,
               mgp = c(2, 0.5, 0))

          # Add grid
          grid(col = "grey80", lty = 1, lwd = 0.5)

          # Redraw points on top of grid
          points(ds_eta[[plot_cols[j]]], ds_eta[[plot_cols[i]]],
                 pch = 16, col = MSDcol[1], cex = 1)

          box()

          # Linear regression line
          abline(lm(ds_eta[[plot_cols[i]]] ~ ds_eta[[plot_cols[j]]]),
                 col = MSDcol[3],
                 lwd = 1.5)
        }
      }
    }

    # Add variable labels on TOP
    for(j in 1:n) {
      mtext(plot_cols[j],
            side = 3, line = 0.2, outer = TRUE,
            at = (j - 0.5)/n, cex = 1, font = 2)
    }

    # Add variable labels on LEFT
    for(i in 1:n) {
      mtext(plot_cols[i],
            side = 2, line = 0.2, outer = TRUE,
            at = (n - i + 0.5)/n, cex = 1, font = 2)
    }

    # Capture the plot
    p_i <- recordPlot()

    # Restore graphics parameters
    par(old_par)

    return(p_i)
  }



  ### Calculate shrinkage
  if (!"ETAshrinkage_var" %in% colnames(smrg_obj$SUMTAB)) {
    message("Shrinkage was not calculated in GFO object")
    smrg_obj$SUMTAB$ETAshrinkage_var <- NA
  }
  shr_out <- smrg_obj$SUMTAB %>% filter(grepl("^omega_", PAR)) %>%
    select(any_of(c("ETAshrinkage_var","VALUE","PAR"))) %>%
    mutate(PAR = str_remove(PAR, "^omega_"), Shrinkage = signif(ETAshrinkage_var,3) )



  ind_par_dist <- patab_i %>% select(all_of(par_seq)) %>%
    gather("PAR") %>%
    left_join(shr_out, by = "PAR") %>%
    mutate(
      Shrinkage = ifelse(is.na(Shrinkage), "N/A", Shrinkage),
      LABEL = str_c(PAR, "\nShrinkage = ", Shrinkage, "%"))


  p_i <- ggplot(data = ind_par_dist, aes(x = value)) +
    geom_histogram(aes(y = after_stat(density)),
                   bins = n_bins,
                   col = "grey25",
                   fill = MSDcol[2]) +
    theme_bw()+
    facet_wrap(~PAR, scales = "free") +
    scale_y_continuous(name = "Density", breaks = scales::pretty_breaks(7), expand = c(0, 0), limits = c(0, NA)) +
    scale_x_continuous(name = "Parameter value", breaks = scales::pretty_breaks(7), expand = c(0, 0))

  suppressMessages({

    if(tdist && plot_type == 'DIST'){

      if (par_type == "Ind") {
        sumtab_i <- smrg_obj$SUMTAB
        if (dplyr::is.grouped_df(sumtab_i)) sumtab_i <- dplyr::ungroup(sumtab_i)
        if ("Distribution" %in% names(sumtab_i) && !"DISTRIBUTION" %in% names(sumtab_i)) {
          sumtab_i <- sumtab_i %>% dplyr::rename(DISTRIBUTION = Distribution)
        }

        dist_tbl <- sumtab_i %>%
          dplyr::filter(str_detect(PAR, "_pop$")) %>%
          dplyr::mutate(PAR = str_remove(PAR, "_pop$")) %>%
          dplyr::select(PAR, dplyr::any_of("DISTRIBUTION"))

        par_pop <- sumtab_i %>%
          dplyr::filter(str_detect(PAR, "_pop$")) %>%
          dplyr::mutate(PAR = str_remove(PAR, "_pop$")) %>%
          dplyr::select(PAR, TV = VALUE) %>%
          dplyr::filter(PAR %in% par_seq) %>%
          dplyr::left_join(dist_tbl, by = "PAR")

        if (!"DISTRIBUTION" %in% names(par_pop)) {
          par_pop_valid <- par_pop[FALSE, , drop = FALSE]
        } else {
          par_pop_valid <- par_pop %>%
            dplyr::filter(
              !is.na(DISTRIBUTION),
              nzchar(trimws(as.character(DISTRIBUTION)))
            )
        }

        if (nrow(par_pop_valid) > 0) {
          colnames(omegamat_i) <- sub("^omega_", "", colnames(omegamat_i))
          rownames(omegamat_i) <- sub("^omega_", "", rownames(omegamat_i))
          rownames(omegamat_i) <- colnames(omegamat_i)
          omega_subset <- omegamat_i[par_pop_valid$PAR, par_pop_valid$PAR, drop = FALSE]

          n_eta <- ncol(omega_subset)

          gen_eta <- MASS::mvrnorm(
            n = 1000,
            mu = rep(0, n_eta),
            Sigma = omega_subset
          ) %>%
            as_tibble() %>%
            setNames(par_pop_valid$PAR) %>%
            tidyr::pivot_longer(cols = dplyr::all_of(par_pop_valid$PAR), names_to = "PAR", values_to = "ETA") %>%
            dplyr::left_join(
              par_pop_valid %>% dplyr::select(PAR, TV, DISTRIBUTION),
              by = "PAR"
            ) %>%
            dplyr::mutate(
              dist_l = tolower(trimws(as.character(DISTRIBUTION))),
              TV_b = pmin(pmax(TV, .Machine$double.eps), 1 - .Machine$double.eps),
              TDIST = dplyr::case_when(
                dist_l == "lognormal" ~ TV * exp(ETA),
                dist_l == "normal" ~ TV + ETA,
                dist_l == "logitnormal" ~ plogis(qlogis(TV_b) + ETA),
                TRUE ~ NA_real_
              )
            ) %>%
            dplyr::filter(!is.na(TDIST)) %>%
            dplyr::left_join(ind_par_dist %>% dplyr::select(PAR, LABEL) %>% dplyr::distinct(), by = "PAR")

          if (nrow(gen_eta) > 0) {
            p_i <- p_i + geom_density(
              data = gen_eta,
              aes(x = TDIST),
              linewidth = 0.8,
              lty = "dashed") +
              facet_wrap(~LABEL, scales = "free")
          }
        }

      } else{
        par_pop <- smrg_obj$SUMTAB %>%
          filter(str_detect(PAR, "^omega_")) %>%
          mutate(PAR_f = str_remove(PAR, "^omega_")) %>%
          filter(PAR_f %in% str_remove(par_seq, "^eta_")) %>%
          select(PAR, TV = VALUE)

        # colnames(omegamat_i) <- sub("^omega_", "", colnames(omegamat_i))
        # rownames(omegamat_i) <- sub("^omega_", "", rownames(omegamat_i))
        rownames(omegamat_i) <- colnames(omegamat_i)
        omega_subset <- omegamat_i[par_pop$PAR, par_pop$PAR, drop = FALSE]

        n_eta <- ncol(omega_subset)

        gen_eta <- MASS::mvrnorm(
          n = 1000,
          mu = rep(0, n_eta),
          Sigma = omega_subset
        ) %>%
          as_tibble() %>%
          setNames(par_pop$PAR) %>%
          pivot_longer(cols = everything(), names_to = "PAR", values_to = "ETA") %>%
          left_join(par_pop, by = "PAR") %>%
          mutate(TDIST =  ETA)%>%
          mutate(PAR = gsub("omega_", "eta_", PAR))

        p_i <- p_i + geom_density(
          data = gen_eta,
          aes(x = TDIST),
          linewidth = 0.8,
          lty = "dashed") +
          facet_wrap(~PAR, scales = "free")
      }









    }


    if (plot_type == 'QQ') {
      eta_cols_all <- names(patab_i)[grepl("^eta_", names(patab_i))]

      if (is.null(par_seq)) {
        eta_selected <- eta_cols_all
      } else {
        # normalize to eta_*
        par_seq_eta <- ifelse(grepl("^eta_", par_seq),
                              par_seq,
                              paste0("eta_", par_seq))

        # check existence
        missing_par <- setdiff(par_seq_eta, names(patab_i))
        if (length(missing_par) > 0) {
          warning(paste("Parameters not found in PATAB:", paste(missing_par, collapse = ", ")))
        }

        eta_selected <- intersect(par_seq_eta, eta_cols_all)
      }
      eta_df <- patab_i %>%
        select(ID, all_of(eta_selected)) %>%
        pivot_longer(cols = -ID, names_to = "PAR", values_to = "ETA")

      p_i <- ggplot(eta_df, aes(sample = ETA)) +
        stat_qq(size = 1.75, color = MSDcol[1], alpha = 0.8) +
        stat_qq_line(col = MSDcol[3]) +
        labs(x = "Theoretical quantiles", y = "Sample quantiles") +
        geom_abline(linewidth = 0.5, col = "black", linetype = "dashed") +
        scale_x_continuous(breaks = scales::pretty_breaks(7)) +
        scale_y_continuous(breaks = scales::pretty_breaks(7)) +
        theme_bw() +
        facet_wrap(~PAR, scales = "free")
    }
  })
  return(p_i)
}
