#!/usr/bin/env Rscript

###############################################################################
# Dynamic Factor Copula Multiple-Decrement Mortality Model (v5 final)
# Revised version:
# - unit detection / rescaling
# - richer diagnostics, tables, and plots
# - coherence checks
# - de-duplicated overlapping age columns (e.g. m95 and m95p)
# - improved paper-oriented cause grouping
###############################################################################

rm(list = ls())
options(stringsAsFactors = FALSE, scipen = 999)

required_pkgs <- c(
  "readxl", "openxlsx", "tidyr", "purrr", "stringr",
  "ggplot2", "scales", "MASS", "dplyr"
)

install_if_missing <- function(pkgs) {
  for (p in pkgs) {
    if (!requireNamespace(p, quietly = TRUE)) {
      install.packages(p, repos = "https://cloud.r-project.org")
    }
  }
}
install_if_missing(required_pkgs)

invisible(lapply(required_pkgs, library, character.only = TRUE))

# Keep dplyr verbs explicit / conflict-safe
select <- dplyr::select
filter <- dplyr::filter
mutate <- dplyr::mutate
summarise <- dplyr::summarise
group_by <- dplyr::group_by
ungroup <- dplyr::ungroup
arrange <- dplyr::arrange
slice_head <- dplyr::slice_head
distinct <- dplyr::distinct
pull <- dplyr::pull
left_join <- dplyr::left_join
bind_rows <- dplyr::bind_rows
rename <- dplyr::rename
lag <- dplyr::lag
count <- dplyr::count

###############################################################################
# CONFIGURATION
###############################################################################

input_dir  <- "C:/Users/Τάσος/Desktop/εγγραφα μου/Phd/papers/dynamic"
output_dir <- "C:/Users/Τάσος/Desktop/εγγραφα μου/Phd/papers/dynamic/data"

auto_detect_file <- TRUE
# target_file <- file.path(input_dir, "your_file.xlsx")

selected_country <- "USA"
selected_sexes   <- c(1, 2)
all_cause_code   <- "L000"

forecast_end_year <- 2070
n_sims           <- 500
seed_value       <- 1234

use_manual_groups <- TRUE
n_major_causes    <- 6

Lseq <- function(a, b) sprintf("L%03d", seq.int(a, b))
manual_groups <- list(
  "Infectious and parasitic diseases" = Lseq(1, 46),
  "Neoplasms" = Lseq(47, 80),
  "Other non-communicable diseases" = Lseq(81, 82),
  "Diabetes and metabolic diseases" = Lseq(83, 86),
  "Mental and neurological disorders" = Lseq(87, 101),
  "Cardiovascular diseases" = Lseq(102, 120),
  "Respiratory diseases" = Lseq(121, 138),
  "Digestive diseases" = Lseq(139, 152),
  "Genitourinary and musculoskeletal diseases" = Lseq(153, 166),
  "Maternal and perinatal conditions" = Lseq(167, 182),
  "Congenital conditions" = Lseq(183, 191),
  "External causes" = Lseq(192, 206)
)

age_col_prefix <- "m"
EPS <- 1e-10

plot_age_max <- 85
selected_plot_ages <- c(0, 1, 40, 65, 80, 85)
selected_profile_years_n <- 5

# Holdout validation configuration
holdout_years <- 10
holdout_min_train_years <- 20
holdout_plot_ages <- c(0, 40, 65, 80)
holdout_age_max <- 85
extra_3d_groups_n <- 3

###############################################################################
# HELPERS
###############################################################################

make_dir <- function(x) {
  if (!dir.exists(x)) dir.create(x, recursive = TRUE, showWarnings = FALSE)
}
make_dir(output_dir)
make_dir(file.path(output_dir, "plots"))
make_dir(file.path(output_dir, "tables"))
make_dir(file.path(output_dir, "rds"))

log_message <- function(...) cat(sprintf("[%s] ", format(Sys.time(), "%Y-%m-%d %H:%M:%S")), sprintf(...), "\n")

sex_label <- function(x) {
  map <- c("0" = "Both sexes", "1" = "Males", "2" = "Females")
  unname(ifelse(as.character(x) %in% names(map), map[as.character(x)], paste("Sex", x)))
}

label_sex_df <- function(df) {
  if (!"sex" %in% names(df)) return(df)
  df$sex_label <- factor(sex_label(df$sex), levels = c("Males", "Females", "Both sexes"))
  df
}

auto_year_grid <- function(years_vec, n = selected_profile_years_n) {
  yrs <- sort(unique(as.integer(years_vec)))
  if (length(yrs) <= n) return(yrs)
  unique(as.integer(round(seq(min(yrs), max(yrs), length.out = n))))
}

save_persp_png <- function(df, x_col, y_col, z_col, filename, main_title, xlab, ylab, zlab, theta = 40, phi = 25) {
  cols_needed <- c(x_col, y_col, z_col)
  tmp <- as.data.frame(df[, cols_needed, drop = FALSE])
  names(tmp) <- c("x", "y", "z")
  tmp$x <- suppressWarnings(as.numeric(tmp$x))
  tmp$y <- suppressWarnings(as.numeric(tmp$y))
  tmp$z <- suppressWarnings(as.numeric(tmp$z))
  tmp <- tmp[stats::complete.cases(tmp), , drop = FALSE]
  if (nrow(tmp) == 0) return(invisible(NULL))

  tmp <- stats::aggregate(z ~ x + y, data = tmp, FUN = mean)
  x_vals <- sort(unique(tmp$x))
  y_vals <- sort(unique(tmp$y))
  if (length(x_vals) < 2 || length(y_vals) < 2) return(invisible(NULL))

  full_grid <- expand.grid(x = x_vals, y = y_vals, KEEP.OUT.ATTRS = FALSE, stringsAsFactors = FALSE)
  tmp2 <- merge(full_grid, tmp, by = c("x", "y"), all.x = TRUE, sort = TRUE)
  tmp2 <- tmp2[order(tmp2$x, tmp2$y), , drop = FALSE]

  # persp() requires nrow(z) == length(x) and ncol(z) == length(y)
  z_mat <- matrix(tmp2$z, nrow = length(x_vals), ncol = length(y_vals), byrow = TRUE)

  # Fill gaps by interpolation along y within each x row; if a full row is missing, borrow nearest valid row
  for (i in seq_len(nrow(z_mat))) {
    row_i <- z_mat[i, ]
    if (all(is.na(row_i))) next
    idx <- which(!is.na(row_i))
    z_mat[i, ] <- stats::approx(x = idx, y = row_i[idx], xout = seq_along(row_i), rule = 2, method = "linear")$y
  }
  if (anyNA(z_mat)) {
    valid_rows <- which(rowSums(is.finite(z_mat)) > 0)
    if (length(valid_rows) == 0) return(invisible(NULL))
    for (i in seq_len(nrow(z_mat))) {
      if (all(!is.finite(z_mat[i, ]))) {
        nearest <- valid_rows[which.min(abs(valid_rows - i))]
        z_mat[i, ] <- z_mat[nearest, ]
      }
    }
  }

  z_mat <- as.matrix(z_mat)
  storage.mode(z_mat) <- "double"
  z_mat[!is.finite(z_mat)] <- 0
  if (!is.matrix(z_mat) || nrow(z_mat) != length(x_vals) || ncol(z_mat) != length(y_vals)) return(invisible(NULL))

  png(filename, width = 1400, height = 1000, res = 160)
  op <- par(no.readonly = TRUE)
  on.exit({par(op); dev.off()}, add = TRUE)
  tryCatch({
    persp(x = x_vals, y = y_vals, z = z_mat, theta = theta, phi = phi, expand = 0.6,
          col = "lightblue", shade = 0.35, ticktype = "detailed",
          xlab = xlab, ylab = ylab, zlab = zlab, main = main_title)
  }, error = function(e) {
    plot.new()
    title(main = paste(main_title, "(surface unavailable)"))
    text(0.5, 0.5, labels = paste("3D surface could not be rendered:", conditionMessage(e)), cex = 0.9)
  })
}

safe_read_file <- function(path) {
  ext <- tolower(tools::file_ext(path))
  if (ext %in% c("xlsx", "xls")) {
    df <- readxl::read_excel(path)
  } else if (ext == "csv") {
    df <- read.csv(path, check.names = FALSE)
  } else {
    stop("Unsupported file extension: ", ext)
  }
  as.data.frame(df)
}

find_candidate_file <- function(folder) {
  files <- list.files(folder, pattern = "\\.(xlsx|xls|csv)$", full.names = TRUE, ignore.case = TRUE)
  if (length(files) == 0) stop("No Excel/CSV files found in: ", folder)
  required_cols <- c("country", "year", "sex", "cause")
  for (f in files) {
    ok <- FALSE
    df <- tryCatch(safe_read_file(f), error = function(e) NULL)
    if (!is.null(df)) {
      nms <- names(df)
      if (all(required_cols %in% nms) && any(grepl(paste0("^", age_col_prefix), nms))) ok <- TRUE
    }
    if (ok) return(f)
  }
  stop("No file with the expected mortality structure was found in: ", folder)
}



resolve_age_columns <- function(age_columns) {
  age_columns <- unique(age_columns)
  if (length(age_columns) == 0) return(age_columns)
  base_num <- suppressWarnings(as.integer(gsub("[^0-9]", "", age_columns)))
  has_p <- grepl("p$", age_columns, ignore.case = TRUE)
  keep <- rep(TRUE, length(age_columns))

  for (b in unique(base_num[!is.na(base_num)])) {
    idx <- which(base_num == b)
    if (length(idx) <= 1) next
    idx_plain <- idx[!has_p[idx]]
    idx_open  <- idx[has_p[idx]]
    if (length(idx_plain) >= 1 && length(idx_open) >= 1) {
      keep[idx_open] <- FALSE
    }
  }

  kept <- age_columns[keep]
  ord <- order(vapply(kept, age_midpoint, numeric(1)), kept)
  kept[ord]
}
age_midpoint <- function(age_label) {
  x <- gsub("^m", "", age_label)
  if (x == "0") return(0)
  if (x == "1") return(1)
  if (x %in% c("95p", "100p", "110p")) {
    base <- suppressWarnings(as.numeric(gsub("p", "", x)))
    return(base)
  }
  suppressWarnings(as.numeric(x))
}

rw_drift_forecast <- function(x, h) {
  x <- as.numeric(x)
  if (length(x) < 3) stop("Time series too short for RW with drift")
  d <- diff(x)
  drift <- mean(d, na.rm = TRUE)
  sigma <- stats::sd(d, na.rm = TRUE)
  mean_fc <- x[length(x)] + drift * seq_len(h)
  list(mean = mean_fc, drift = drift, sigma = ifelse(is.na(sigma), 0, sigma))
}

fit_lc_svd <- function(rate_matrix) {
  mat <- as.matrix(rate_matrix)
  storage.mode(mat) <- "double"
  mat[!is.finite(mat)] <- NA_real_
  mat <- pmax(mat, EPS)
  logm <- log(mat)

  ax <- rowMeans(logm, na.rm = TRUE)
  centered <- sweep(logm, 1, ax, FUN = "-")
  centered[is.na(centered)] <- 0

  sv <- svd(centered)
  bx <- sv$u[, 1]
  kt <- sv$d[1] * sv$v[, 1]

  bx <- bx / sum(bx)
  kt <- kt * sum(sv$u[, 1])

  fitted_logm <- outer(ax, rep(1, length(kt))) + bx %o% kt
  fitted_m <- exp(fitted_logm)

  list(ax = ax, bx = bx, kt = kt, fitted = fitted_m, log_obs = logm, log_fit = fitted_logm)
}

simulate_gaussian_copula_innovations <- function(corr_mat, n_sims, h, std_vec) {
  p <- ncol(corr_mat)
  if (p == 1) {
    arr <- array(rnorm(n_sims * h, mean = 0, sd = std_vec[1]), dim = c(n_sims, h, 1))
    return(arr)
  }
  corr_mat <- as.matrix(corr_mat)
  diag(corr_mat) <- 1
  eig <- eigen(corr_mat, symmetric = TRUE)
  eig$values[eig$values < 1e-8] <- 1e-8
  corr_pd <- eig$vectors %*% diag(eig$values) %*% t(eig$vectors)
  D <- diag(sqrt(diag(corr_pd)))
  corr_pd <- solve(D) %*% corr_pd %*% solve(D)

  out <- array(NA_real_, dim = c(n_sims, h, p))
  for (tt in seq_len(h)) {
    z <- MASS::mvrnorm(n = n_sims, mu = rep(0, p), Sigma = corr_pd)
    z <- sweep(z, 2, std_vec, `*`)
    out[, tt, ] <- z
  }
  out
}

normalize_rows <- function(mat) {
  rs <- rowSums(mat)
  rs[rs <= 0 | !is.finite(rs)] <- 1
  mat / rs
}

infer_interval_width <- function(age_values) {
  age_values <- sort(unique(age_values))
  if (length(age_values) == 1) {
    return(data.frame(age = age_values, n = 1))
  }
  widths <- c(diff(age_values), tail(diff(age_values), 1))
  widths[!is.finite(widths) | widths <= 0] <- 5
  data.frame(age = age_values, n = widths)
}

safe_label <- function(x) {
  gsub("[^A-Za-z0-9_\\-]+", "_", x)
}

save_plot <- function(plot_obj, filename, width = 11, height = 6.5) {
  ggplot2::ggsave(filename, plot_obj, width = width, height = height, dpi = 300)
}

compute_accuracy_metrics <- function(actual, forecast) {
  actual <- as.numeric(actual)
  forecast <- as.numeric(forecast)
  err <- forecast - actual
  ape <- abs(err) / pmax(abs(actual), EPS)
  data.frame(
    MAE = mean(abs(err), na.rm = TRUE),
    RMSE = sqrt(mean(err^2, na.rm = TRUE)),
    MAPE = mean(ape, na.rm = TRUE),
    bias = mean(err, na.rm = TRUE)
  )
}

fit_lc_svd_from_train_to_forecast <- function(df_train, df_test_years) {
  mat <- df_train %>%
    dplyr::select(age, year, mx) %>%
    tidyr::pivot_wider(names_from = year, values_from = mx, values_fn = mean) %>%
    arrange(age)

  age_vec <- mat$age
  rate_mat <- as.matrix(mat[, -1, drop = FALSE])
  storage.mode(rate_mat) <- "double"
  col_years <- as.integer(colnames(rate_mat))

  lc <- fit_lc_svd(rate_mat)
  h <- length(df_test_years)
  kt_fc <- rw_drift_forecast(lc$kt, h)
  mean_logm_fc <- outer(lc$ax, rep(1, h)) + lc$bx %o% kt_fc$mean
  data.frame(
    age = rep(age_vec, h),
    year = rep(df_test_years, each = length(age_vec)),
    mx_forecast = as.vector(exp(mean_logm_fc))
  )
}

###############################################################################
# DATA INGESTION
###############################################################################

if (auto_detect_file) {
  target_file <- find_candidate_file(input_dir)
}
log_message("Using source file: %s", target_file)

raw <- safe_read_file(target_file)
names(raw) <- trimws(names(raw))

required_cols <- c("country", "year", "sex", "cause")
missing_cols <- setdiff(required_cols, names(raw))
if (length(missing_cols) > 0) stop("Missing required columns: ", paste(missing_cols, collapse = ", "))

age_cols <- names(raw)[grepl(paste0("^", age_col_prefix), names(raw))]
age_cols <- resolve_age_columns(age_cols)
if (length(age_cols) == 0) stop("No age-specific mortality columns found, e.g. m0, m1, m5, ...")
log_message("Age columns retained for modeling: %s", paste(age_cols, collapse = ", "))

raw <- raw %>%
  mutate(
    year = as.integer(year),
    sex  = as.integer(sex),
    cause = as.character(cause),
    country = as.character(country)
  )

mort_long_raw <- raw %>%
  dplyr::filter(country == selected_country, sex %in% selected_sexes) %>%
  tidyr::pivot_longer(cols = all_of(age_cols), names_to = "age_label", values_to = "mx_raw") %>%
  mutate(
    age = vapply(age_label, age_midpoint, numeric(1)),
    mx_raw = as.numeric(mx_raw)
  ) %>%
  dplyr::filter(!is.na(age), !is.na(mx_raw), is.finite(mx_raw), mx_raw >= 0)

if (nrow(mort_long_raw) == 0) stop("No data remain after filtering. Check country/sex codes.")

mort_long_raw <- mort_long_raw %>%
  group_by(country, year, sex, cause, age, age_label) %>%
  summarise(mx_raw = mean(mx_raw, na.rm = TRUE), .groups = "drop")

allcause_raw <- mort_long_raw %>%
  filter(cause == all_cause_code) %>%
  group_by(country, year, sex, age, age_label) %>%
  summarise(mx_raw = mean(mx_raw, na.rm = TRUE), .groups = "drop")
if (nrow(allcause_raw) == 0) stop("All-cause code not found: ", all_cause_code)

# Unit detection
unit_stats <- allcause_raw %>%
  summarise(
    min_mx = min(mx_raw, na.rm = TRUE),
    p50_mx = stats::median(mx_raw, na.rm = TRUE),
    p95_mx = stats::quantile(mx_raw, 0.95, na.rm = TRUE),
    max_mx = max(mx_raw, na.rm = TRUE)
  )

unit_scale <- 1
unit_inference <- "Rates appear already on probability / rate scale."
if (is.finite(unit_stats$max_mx) && unit_stats$max_mx > 50) {
  unit_scale <- 1e5
  unit_inference <- "Detected very large mortality rates; treating source data as rates per 100,000 and rescaling by 100,000."
} else if (is.finite(unit_stats$max_mx) && unit_stats$max_mx > 5) {
  unit_scale <- 1e3
  unit_inference <- "Detected elevated mortality rates; treating source data as rates per 1,000 and rescaling by 1,000."
}

unit_diagnostics <- data.frame(
  detected_scale_divisor = unit_scale,
  interpretation = unit_inference,
  min_mx_raw = unit_stats$min_mx,
  p50_mx_raw = unit_stats$p50_mx,
  p95_mx_raw = unit_stats$p95_mx,
  max_mx_raw = unit_stats$max_mx,
  stringsAsFactors = FALSE
)

log_message("%s", unit_inference)

mort_long <- mort_long_raw %>%
  mutate(mx = pmax(mx_raw / unit_scale, EPS))
allcause_df <- allcause_raw %>%
  mutate(mx = pmax(mx_raw / unit_scale, EPS))

meta_summary <- mort_long %>%
  summarise(
    n_rows = n(),
    min_year = min(year),
    max_year = max(year),
    n_causes = n_distinct(cause),
    n_sexes = n_distinct(sex),
    n_ages = n_distinct(age),
    unit_scale_divisor = unit_scale
  )

log_message("Rows: %s | Years: %s-%s | Causes: %s | Sexes: %s | Ages: %s",
            meta_summary$n_rows, meta_summary$min_year, meta_summary$max_year,
            meta_summary$n_causes, meta_summary$n_sexes, meta_summary$n_ages)

###############################################################################
# DECREMENT GROUP CONSTRUCTION
###############################################################################

top_causes_table <- NULL

if (use_manual_groups && length(manual_groups) > 0) {
  log_message("Using manual decrement groups.")

  cause_map <- bind_rows(lapply(names(manual_groups), function(g) {
    data.frame(cause = manual_groups[[g]], dec_group = g, stringsAsFactors = FALSE)
  }))

  grouped <- mort_long %>%
    dplyr::filter(cause != all_cause_code) %>%
    left_join(cause_map, by = "cause") %>%
    mutate(dec_group = ifelse(is.na(dec_group), "OTHER", dec_group)) %>%
    group_by(country, year, sex, age, age_label, dec_group) %>%
    summarise(mx = sum(mx, na.rm = TRUE), .groups = "drop")
} else {
  log_message("Using automatic decrement grouping: top %d causes + OTHER.", n_major_causes)

  top_causes_table <- mort_long %>%
    dplyr::filter(cause != all_cause_code) %>%
    group_by(cause) %>%
    summarise(avg_total = mean(mx, na.rm = TRUE), .groups = "drop") %>%
    arrange(desc(avg_total))

  top_causes <- top_causes_table %>%
    slice_head(n = n_major_causes) %>%
    pull(cause)

  grouped <- mort_long %>%
    dplyr::filter(cause != all_cause_code) %>%
    mutate(dec_group = ifelse(cause %in% top_causes, cause, "OTHER")) %>%
    group_by(country, year, sex, age, age_label, dec_group) %>%
    summarise(mx = sum(mx, na.rm = TRUE), .groups = "drop")
}

share_df <- grouped %>%
  left_join(allcause_df %>% dplyr::select(country, year, sex, age, age_label, all_mx = mx),
            by = c("country", "year", "sex", "age", "age_label")) %>%
  mutate(
    all_mx = pmax(all_mx, EPS),
    share = pmin(pmax(mx / all_mx, EPS), 1 - EPS)
  )

valid_groups <- share_df %>%
  group_by(dec_group, sex) %>%
  summarise(n_years = n_distinct(year), .groups = "drop") %>%
  dplyr::filter(n_years >= 10) %>%
  distinct(dec_group) %>%
  pull(dec_group)

share_df <- share_df %>% filter(dec_group %in% valid_groups)

if (!"OTHER" %in% share_df$dec_group) {
  other_df <- allcause_df %>%
    dplyr::select(country, year, sex, age, age_label, all_mx = mx) %>%
    mutate(dec_group = "OTHER", mx = EPS, share = EPS)
  share_df <- bind_rows(share_df, other_df)
}

share_df <- share_df %>%
  group_by(country, year, sex, age, age_label) %>%
  mutate(share = share / sum(share, na.rm = TRUE)) %>%
  ungroup()

selected_groups <- share_df %>% distinct(dec_group) %>% arrange(dec_group) %>% pull(dec_group)
selected_groups_table <- data.frame(dec_group = selected_groups, stringsAsFactors = FALSE)
if (use_manual_groups && length(manual_groups) > 0) {
  selected_groups_table$n_codes <- sapply(selected_groups, function(g) sum(cause_map$dec_group == g))
  selected_groups_table$code_examples <- sapply(selected_groups, function(g) paste(head(cause_map$cause[cause_map$dec_group == g], 5), collapse = ", "))
} else {
  selected_groups_table$n_codes <- NA_integer_
  selected_groups_table$code_examples <- selected_groups
}
log_message("Selected decrement groups: %s", paste(selected_groups, collapse = ", "))

share_sum_check_hist <- share_df %>%
  group_by(year, sex, age) %>%
  summarise(sum_share = sum(share, na.rm = TRUE), .groups = "drop") %>%
  mutate(abs_error = abs(sum_share - 1))

###############################################################################
# ALL-CAUSE LEE-CARTER FIT BY SEX
###############################################################################

years <- sort(unique(allcause_df$year))
ages  <- sort(unique(allcause_df$age))
forecast_horizon <- max(1, forecast_end_year - max(years))
log_message("Forecast end year set to %s -> horizon %s years.", forecast_end_year, forecast_horizon)

lc_results <- list()
allcause_forecasts <- list()
allcause_sim_paths <- list()
kappa_summary <- list()

for (sx in selected_sexes) {
  df_sx <- allcause_df %>%
    filter(sex == sx) %>%
    group_by(age, year) %>%
    summarise(mx = mean(mx, na.rm = TRUE), .groups = "drop")

  dup_check_all <- df_sx %>% count(age, year) %>% filter(n > 1)
  if (nrow(dup_check_all) > 0) stop(sprintf("Duplicate all-cause cells remain for sex=%s after collapsing.", sx))

  mat <- df_sx %>%
    dplyr::select(age, year, mx) %>%
    tidyr::pivot_wider(names_from = year, values_from = mx, values_fn = mean) %>%
    arrange(age)

  age_vec <- mat$age
  rate_mat <- as.matrix(mat[, -1, drop = FALSE])
  storage.mode(rate_mat) <- "double"
  col_years <- as.integer(colnames(rate_mat))

  lc <- fit_lc_svd(rate_mat)
  kt_fc <- rw_drift_forecast(lc$kt, forecast_horizon)

  set.seed(seed_value + sx)
  sim_innov <- matrix(rnorm(n_sims * forecast_horizon, mean = 0, sd = kt_fc$sigma), nrow = n_sims, ncol = forecast_horizon)
  kt_paths <- matrix(NA_real_, nrow = n_sims, ncol = forecast_horizon)
  for (ss in seq_len(n_sims)) {
    current <- tail(lc$kt, 1)
    for (hh in seq_len(forecast_horizon)) {
      current <- current + kt_fc$drift + sim_innov[ss, hh]
      kt_paths[ss, hh] <- current
    }
  }

  fc_years <- max(col_years) + seq_len(forecast_horizon)
  mean_logm_fc <- outer(lc$ax, rep(1, forecast_horizon)) + lc$bx %o% kt_fc$mean
  mean_m_fc <- exp(mean_logm_fc)

  lc_results[[paste0("sex_", sx)]] <- list(
    sex = sx,
    ages = age_vec,
    years = col_years,
    lc = lc,
    kappa_forecast = kt_fc,
    fc_years = fc_years
  )

  allcause_forecasts[[paste0("sex_", sx)]] <-
    expand.grid(age = age_vec, year = fc_years, KEEP.OUT.ATTRS = FALSE, stringsAsFactors = FALSE) %>%
    mutate(
      sex = sx,
      mx_forecast_all = as.vector(mean_m_fc)
    )

  allcause_sim_paths[[paste0("sex_", sx)]] <- list(ages = age_vec, years = fc_years, kt_paths = kt_paths)

  kappa_summary[[paste0("sex_", sx)]] <- data.frame(
    sex = sx,
    year = c(col_years, fc_years),
    kappa = c(lc$kt, kt_fc$mean),
    type = c(rep("Historical", length(col_years)), rep("Forecast", length(fc_years))),
    stringsAsFactors = FALSE
  )
}

allcause_forecast_df <- bind_rows(allcause_forecasts)
kappa_summary_df <- bind_rows(kappa_summary)

###############################################################################
# SHARE MODELS BY SEX AND DECREMENT GROUP
###############################################################################

share_models <- list()
share_forecasts <- list()
correlation_objects <- list()
share_simulations <- list()

for (sx in selected_sexes) {
  sx_data <- share_df %>% filter(sex == sx)
  groups_sx <- sx_data %>% distinct(dec_group) %>% arrange(dec_group) %>% pull(dec_group)
  groups_model <- setdiff(groups_sx, "OTHER")
  if (length(groups_model) == 0) next

  group_kt <- list()
  group_fc <- list()
  innov_table <- list()

  for (g in groups_model) {
    dg <- sx_data %>%
      filter(dec_group == g) %>%
      group_by(age, year) %>%
      summarise(share = mean(share, na.rm = TRUE), .groups = "drop")

    dup_check_share <- dg %>% count(age, year) %>% filter(n > 1)
    if (nrow(dup_check_share) > 0) stop(sprintf("Duplicate share cells remain for sex=%s, group=%s after collapsing.", sx, g))

    mat <- dg %>%
      mutate(logit_share = qlogis(pmin(pmax(share, EPS), 1 - EPS))) %>%
      dplyr::select(age, year, logit_share) %>%
      tidyr::pivot_wider(names_from = year, values_from = logit_share, values_fn = mean) %>%
      arrange(age)

    age_vec <- mat$age
    z_mat <- as.matrix(mat[, -1, drop = FALSE])
    storage.mode(z_mat) <- "double"
    col_years <- as.integer(colnames(z_mat))

    lc_z <- fit_lc_svd(exp(z_mat))
    kt_fc <- rw_drift_forecast(lc_z$kt, forecast_horizon)
    innov <- diff(lc_z$kt)

    innov_table[[g]] <- data.frame(year = col_years[-1], group = g, innov = innov, stringsAsFactors = FALSE)
    mean_z_fc <- outer(lc_z$ax, rep(1, forecast_horizon)) + lc_z$bx %o% kt_fc$mean

    group_kt[[g]] <- list(age = age_vec, years = col_years, model = lc_z, fc = kt_fc)
    group_fc[[g]] <- list(age = age_vec, fc_years = max(col_years) + seq_len(forecast_horizon), z_fc = mean_z_fc)
  }

  innov_df <- bind_rows(innov_table)
  innov_wide <- innov_df %>%
    tidyr::pivot_wider(names_from = group, values_from = innov) %>%
    arrange(year)

  innov_mat <- as.matrix(innov_wide[, setdiff(names(innov_wide), "year"), drop = FALSE])
  storage.mode(innov_mat) <- "double"

  if (ncol(innov_mat) == 1) {
    corr_mat <- matrix(1, nrow = 1, ncol = 1)
    colnames(corr_mat) <- rownames(corr_mat) <- colnames(innov_mat)
  } else {
    corr_mat <- cor(innov_mat, use = "pairwise.complete.obs")
    corr_mat[!is.finite(corr_mat)] <- 0
    diag(corr_mat) <- 1
  }

  sigma_vec <- sapply(groups_model, function(g) group_kt[[g]]$fc$sigma)
  names(sigma_vec) <- groups_model

  set.seed(seed_value + 100 + sx)
  copula_shocks <- simulate_gaussian_copula_innovations(corr_mat = corr_mat, n_sims = n_sims, h = forecast_horizon, std_vec = sigma_vec)

  sim_group_paths <- list()
  for (jj in seq_along(groups_model)) {
    g <- groups_model[jj]
    current <- tail(group_kt[[g]]$model$kt, 1)
    drift_g <- group_kt[[g]]$fc$drift
    kt_paths <- matrix(NA_real_, nrow = n_sims, ncol = forecast_horizon)
    for (ss in seq_len(n_sims)) {
      cur <- current
      for (hh in seq_len(forecast_horizon)) {
        cur <- cur + drift_g + copula_shocks[ss, hh, jj]
        kt_paths[ss, hh] <- cur
      }
    }
    sim_group_paths[[g]] <- kt_paths
  }

  mean_scores_list <- list()
  fc_years <- NULL
  age_template <- NULL
  for (g in groups_model) {
    gf <- group_fc[[g]]
    z_fc <- gf$z_fc
    score <- pmax(exp(z_fc), EPS)
    mean_scores_list[[g]] <- score
    fc_years <- gf$fc_years
    age_template <- gf$age
  }

  score_array <- array(0, dim = c(length(age_template), forecast_horizon, length(groups_sx)),
                       dimnames = list(age = age_template, year = fc_years, group = groups_sx))

  for (g in groups_model) score_array[, , g] <- mean_scores_list[[g]]
  score_array[, , "OTHER"] <- 1

  share_array <- array(NA_real_, dim = dim(score_array), dimnames = dimnames(score_array))
  for (hh in seq_len(forecast_horizon)) {
    share_array[, hh, ] <- normalize_rows(score_array[, hh, ])
  }

  share_forecasts[[paste0("sex_", sx)]] <-
    expand.grid(age = age_template, year = fc_years, dec_group = groups_sx,
                KEEP.OUT.ATTRS = FALSE, stringsAsFactors = FALSE) %>%
    mutate(
      sex = sx,
      share_forecast = as.vector(share_array)
    )

  share_models[[paste0("sex_", sx)]] <- group_kt
  correlation_objects[[paste0("sex_", sx)]] <- list(
    corr_mat = corr_mat,
    innov_df = innov_df
  )
  share_simulations[[paste0("sex_", sx)]] <- list(
    ages = age_template,
    fc_years = fc_years,
    all_groups = groups_sx,
    model_groups = groups_model,
    sim_group_paths = sim_group_paths
  )
}

share_forecast_df <- bind_rows(share_forecasts)

###############################################################################
# COMBINE ALL-CAUSE + SHARES => DECREMENT FORECASTS
###############################################################################

forecast_rates_df <- share_forecast_df %>%
  left_join(allcause_forecast_df, by = c("sex", "age", "year")) %>%
  mutate(mx_forecast_dec = pmax(mx_forecast_all, EPS) * pmax(share_forecast, EPS))

coherence_forecast_df <- forecast_rates_df %>%
  group_by(year, sex, age) %>%
  summarise(sum_dec_mx = sum(mx_forecast_dec, na.rm = TRUE), .groups = "drop") %>%
  left_join(allcause_forecast_df %>% rename(allcause_mx = mx_forecast_all), by = c("year", "sex", "age")) %>%
  mutate(abs_error = abs(sum_dec_mx - allcause_mx), rel_error = abs_error / pmax(allcause_mx, EPS))

###############################################################################
# SIMULATION OUTPUTS FOR PREDICTIVE DISTRIBUTIONS
###############################################################################

simulation_summary_list <- list()
simulation_allcause_list <- list()

for (sx in selected_sexes) {
  key <- paste0("sex_", sx)
  if (!key %in% names(share_simulations)) next

  all_obj <- allcause_sim_paths[[key]]
  share_obj <- share_simulations[[key]]
  lc_obj <- lc_results[[key]]

  ages_fc <- share_obj$ages
  years_fc <- share_obj$fc_years
  groups_sx <- share_obj$all_groups
  groups_model <- names(share_obj$sim_group_paths)

  all_rates_arr <- array(NA_real_, dim = c(n_sims, length(ages_fc), forecast_horizon))
  for (ss in seq_len(n_sims)) {
    logm <- outer(lc_obj$lc$ax, rep(1, forecast_horizon)) + lc_obj$lc$bx %o% all_obj$kt_paths[ss, ]
    all_rates_arr[ss, , ] <- exp(logm)
  }

  med_all <- apply(all_rates_arr, c(2, 3), median, na.rm = TRUE)
  lo_all  <- apply(all_rates_arr, c(2, 3), quantile, probs = 0.025, na.rm = TRUE)
  hi_all  <- apply(all_rates_arr, c(2, 3), quantile, probs = 0.975, na.rm = TRUE)

  simulation_allcause_list[[key]] <- expand.grid(age = ages_fc, year = years_fc, KEEP.OUT.ATTRS = FALSE, stringsAsFactors = FALSE) %>%
    mutate(
      sex = sx,
      median_mx_all = as.vector(med_all),
      low95_mx_all = as.vector(lo_all),
      high95_mx_all = as.vector(hi_all)
    )

  dec_arr <- array(NA_real_, dim = c(n_sims, length(ages_fc), forecast_horizon, length(groups_sx)),
                   dimnames = list(sim = NULL, age = ages_fc, year = years_fc, group = groups_sx))

  for (ss in seq_len(n_sims)) {
    score_arr <- array(0, dim = c(length(ages_fc), forecast_horizon, length(groups_sx)),
                       dimnames = list(age = ages_fc, year = years_fc, group = groups_sx))

    for (g in groups_model) {
      g_model <- share_models[[key]][[g]]$model
      kt_path <- share_obj$sim_group_paths[[g]][ss, ]
      z_future <- outer(g_model$ax, rep(1, forecast_horizon)) + g_model$bx %o% kt_path
      score_arr[, , g] <- pmax(exp(z_future), EPS)
    }
    score_arr[, , "OTHER"] <- 1

    for (hh in seq_len(forecast_horizon)) {
      share_h <- normalize_rows(score_arr[, hh, ])
      for (gg in seq_along(groups_sx)) {
        dec_arr[ss, , hh, gg] <- all_rates_arr[ss, , hh] * share_h[, gg]
      }
    }
  }

  for (gg in seq_along(groups_sx)) {
    g <- groups_sx[gg]
    med <- apply(dec_arr[, , , gg, drop = FALSE], c(2, 3), median, na.rm = TRUE)
    lo  <- apply(dec_arr[, , , gg, drop = FALSE], c(2, 3), quantile, probs = 0.025, na.rm = TRUE)
    hi  <- apply(dec_arr[, , , gg, drop = FALSE], c(2, 3), quantile, probs = 0.975, na.rm = TRUE)

    simulation_summary_list[[paste0(key, "_", g)]] <-
      expand.grid(age = ages_fc, year = years_fc, KEEP.OUT.ATTRS = FALSE, stringsAsFactors = FALSE) %>%
      mutate(
        sex = sx,
        dec_group = g,
        median_mx = as.vector(med),
        low95_mx  = as.vector(lo),
        high95_mx = as.vector(hi)
      )
  }
}

simulation_summary_df <- bind_rows(simulation_summary_list)
simulation_allcause_df <- bind_rows(simulation_allcause_list)

###############################################################################
# LIFE TABLE STYLE DECREMENT PROBABILITIES
###############################################################################

age_width_df <- infer_interval_width(sort(unique(forecast_rates_df$age)))

decrement_q_df <- forecast_rates_df %>%
  left_join(age_width_df, by = "age") %>%
  mutate(
    q_forecast_dec = 1 - exp(-pmax(mx_forecast_dec, 0) * n),
    q_forecast_dec = pmin(pmax(q_forecast_dec, 0), 1)
  )

allcause_q_df <- allcause_forecast_df %>%
  left_join(age_width_df, by = "age") %>%
  mutate(
    q_forecast_all = 1 - exp(-pmax(mx_forecast_all, 0) * n),
    q_forecast_all = pmin(pmax(q_forecast_all, 0), 1)
  )

qx_summary <- decrement_q_df %>%
  summarise(
    min_q = min(q_forecast_dec, na.rm = TRUE),
    p50_q = median(q_forecast_dec, na.rm = TRUE),
    p95_q = quantile(q_forecast_dec, 0.95, na.rm = TRUE),
    max_q = max(q_forecast_dec, na.rm = TRUE)
  )

###############################################################################
# DIAGNOSTIC TABLES
###############################################################################

sex_labels <- data.frame(sex = selected_sexes, sex_label = paste0("sex_", selected_sexes), stringsAsFactors = FALSE)

historical_age_year_table <- allcause_df %>%
  group_by(year, sex, age) %>%
  summarise(mx = mean(mx, na.rm = TRUE), .groups = "drop") %>%
  filter(age %in% selected_plot_ages)

forecast_age_year_table <- allcause_forecast_df %>%
  filter(age %in% selected_plot_ages)

forecast_summary_table <- forecast_rates_df %>%
  group_by(sex, dec_group, age) %>%
  summarise(
    start_year = min(year),
    end_year = max(year),
    start_mx = mx_forecast_dec[year == min(year)][1],
    end_mx = mx_forecast_dec[year == max(year)][1],
    pct_change = (end_mx / pmax(start_mx, EPS)) - 1,
    mean_share = mean(share_forecast, na.rm = TRUE),
    .groups = "drop"
  )

coherence_summary <- coherence_forecast_df %>%
  summarise(
    max_abs_error = max(abs_error, na.rm = TRUE),
    mean_abs_error = mean(abs_error, na.rm = TRUE),
    max_rel_error = max(rel_error, na.rm = TRUE),
    mean_rel_error = mean(rel_error, na.rm = TRUE)
  )

correlation_table <- bind_rows(lapply(names(correlation_objects), function(nm) {
  obj <- correlation_objects[[nm]]
  cm <- obj$corr_mat
  if (is.null(cm)) return(NULL)
  as.data.frame(as.table(cm), stringsAsFactors = FALSE) %>%
    rename(group1 = Var1, group2 = Var2, correlation = Freq) %>%
    mutate(sex = as.integer(gsub("sex_", "", nm)), sex_label = factor(sex_label(as.integer(gsub("sex_", "", nm))), levels = c("Males", "Females", "Both sexes")))
}))

interval_table <- age_width_df

allcause_percentiles <- allcause_forecast_df %>%
  group_by(sex, age) %>%
  summarise(
    min_fc = min(mx_forecast_all, na.rm = TRUE),
    med_fc = median(mx_forecast_all, na.rm = TRUE),
    max_fc = max(mx_forecast_all, na.rm = TRUE),
    .groups = "drop"
  )

share_forecast_check <- share_forecast_df %>%
  group_by(year, sex, age) %>%
  summarise(sum_share_forecast = sum(share_forecast, na.rm = TRUE), .groups = "drop") %>%
  mutate(abs_error = abs(sum_share_forecast - 1))

allcause_df <- label_sex_df(allcause_df)
allcause_forecast_df <- label_sex_df(allcause_forecast_df)
kappa_summary_df <- label_sex_df(kappa_summary_df)
share_forecast_df <- label_sex_df(share_forecast_df)
forecast_rates_df <- label_sex_df(forecast_rates_df)
decrement_q_df <- label_sex_df(decrement_q_df)
simulation_summary_df <- label_sex_df(simulation_summary_df)
simulation_allcause_df <- label_sex_df(simulation_allcause_df)
coherence_forecast_df <- label_sex_df(coherence_forecast_df)
allcause_q_df <- label_sex_df(allcause_q_df)
correlation_table <- label_sex_df(correlation_table)
share_sum_check_hist <- label_sex_df(share_sum_check_hist)
share_forecast_check <- label_sex_df(share_forecast_check)

###############################################################################
# HOLDOUT VALIDATION FOR ALL-CAUSE BACKBONE
###############################################################################

holdout_results_df <- data.frame()
holdout_metrics_by_age <- data.frame()
holdout_metrics_overall <- data.frame()
holdout_error_surface_df <- data.frame()

for (sx in selected_sexes) {
  df_sx <- allcause_df %>%
    filter(sex == sx) %>%
    group_by(age, year) %>%
    summarise(mx = mean(mx, na.rm = TRUE), .groups = "drop")

  years_sx <- sort(unique(df_sx$year))
  if (length(years_sx) <= (holdout_years + holdout_min_train_years)) next

  test_years <- tail(years_sx, holdout_years)
  train_years <- setdiff(years_sx, test_years)
  if (length(train_years) < holdout_min_train_years) next

  df_train <- df_sx %>% filter(year %in% train_years)
  df_test  <- df_sx %>% filter(year %in% test_years)
  df_fc    <- fit_lc_svd_from_train_to_forecast(df_train, test_years)

  df_hold <- df_test %>%
    left_join(df_fc, by = c("age", "year")) %>%
    mutate(
      sex = sx,
      abs_error = abs(mx_forecast - mx),
      pct_error = (mx_forecast - mx) / pmax(mx, EPS),
      ape = abs(mx_forecast - mx) / pmax(abs(mx), EPS)
    )

  holdout_results_df <- bind_rows(holdout_results_df, df_hold)

  met_age <- df_hold %>%
    group_by(sex, age) %>%
    summarise(
      MAE = mean(abs_error, na.rm = TRUE),
      RMSE = sqrt(mean((mx_forecast - mx)^2, na.rm = TRUE)),
      MAPE = mean(ape, na.rm = TRUE),
      bias = mean(mx_forecast - mx, na.rm = TRUE),
      .groups = "drop"
    )
  holdout_metrics_by_age <- bind_rows(holdout_metrics_by_age, met_age)

  met_overall <- df_hold %>%
    summarise(
      sex = sx,
      train_start = min(train_years),
      train_end = max(train_years),
      test_start = min(test_years),
      test_end = max(test_years),
      MAE = mean(abs_error, na.rm = TRUE),
      RMSE = sqrt(mean((mx_forecast - mx)^2, na.rm = TRUE)),
      MAPE = mean(ape, na.rm = TRUE),
      bias = mean(mx_forecast - mx, na.rm = TRUE)
    )
  holdout_metrics_overall <- bind_rows(holdout_metrics_overall, met_overall)
}

if (nrow(holdout_results_df) > 0) {
  holdout_results_df <- label_sex_df(holdout_results_df)
  holdout_metrics_by_age <- label_sex_df(holdout_metrics_by_age)
  holdout_metrics_overall <- label_sex_df(holdout_metrics_overall)
  holdout_error_surface_df <- holdout_results_df %>%
    mutate(log_abs_error = log10(pmax(abs_error, EPS)),
           abs_pct_error = abs(pct_error))
}

###############################################################################
# OUTPUT TABLES
###############################################################################

log_message("Writing tables...")

wb <- openxlsx::createWorkbook()

write_sheet <- function(wb, sheet, df) {
  openxlsx::addWorksheet(wb, sheet)
  openxlsx::writeData(wb, sheet, df)
}

write_sheet(wb, "meta_summary", meta_summary)
write_sheet(wb, "unit_diagnostics", unit_diagnostics)
write_sheet(wb, "selected_groups", selected_groups_table)
if (!is.null(top_causes_table)) write_sheet(wb, "top_causes", top_causes_table)
if (exists("cause_map")) write_sheet(wb, "group_definition", cause_map)
write_sheet(wb, "age_widths", interval_table)
write_sheet(wb, "hist_share_sum_check", share_sum_check_hist)
write_sheet(wb, "fc_share_sum_check", share_forecast_check)
write_sheet(wb, "coherence_forecast", coherence_forecast_df)
write_sheet(wb, "coherence_summary", coherence_summary)
write_sheet(wb, "kappa_summary", kappa_summary_df)
write_sheet(wb, "allcause_forecast", allcause_forecast_df)
write_sheet(wb, "allcause_q", allcause_q_df)
write_sheet(wb, "share_forecast", share_forecast_df)
write_sheet(wb, "decrement_rates", forecast_rates_df)
write_sheet(wb, "decrement_q", decrement_q_df)
write_sheet(wb, "sim_allcause", simulation_allcause_df)
write_sheet(wb, "sim_decrement", simulation_summary_df)
write_sheet(wb, "forecast_summary", forecast_summary_table)
write_sheet(wb, "correlation_table", correlation_table)
write_sheet(wb, "allcause_percentiles", allcause_percentiles)
write_sheet(wb, "qx_summary", qx_summary)
write_sheet(wb, "hist_age_year", historical_age_year_table)
write_sheet(wb, "fc_age_year", forecast_age_year_table)
if (nrow(holdout_results_df) > 0) write_sheet(wb, "holdout_allcause", holdout_results_df)
if (nrow(holdout_metrics_by_age) > 0) write_sheet(wb, "holdout_metrics_age", holdout_metrics_by_age)
if (nrow(holdout_metrics_overall) > 0) write_sheet(wb, "holdout_metrics_overall", holdout_metrics_overall)

results_xlsx <- file.path(output_dir, "tables", "dynamic_factor_copula_mortality_results_v5.xlsx")
openxlsx::saveWorkbook(wb, results_xlsx, overwrite = TRUE)

write.csv(unit_diagnostics, file.path(output_dir, "tables", "unit_diagnostics.csv"), row.names = FALSE)
if (!is.null(top_causes_table)) write.csv(top_causes_table, file.path(output_dir, "tables", "top_causes.csv"), row.names = FALSE)
write.csv(allcause_forecast_df, file.path(output_dir, "tables", "allcause_forecast.csv"), row.names = FALSE)
write.csv(allcause_q_df, file.path(output_dir, "tables", "allcause_q_forecast.csv"), row.names = FALSE)
write.csv(kappa_summary_df, file.path(output_dir, "tables", "kappa_summary.csv"), row.names = FALSE)
write.csv(share_forecast_df, file.path(output_dir, "tables", "share_forecast.csv"), row.names = FALSE)
write.csv(forecast_rates_df, file.path(output_dir, "tables", "decrement_rate_forecast.csv"), row.names = FALSE)
write.csv(decrement_q_df, file.path(output_dir, "tables", "decrement_q_forecast.csv"), row.names = FALSE)
write.csv(simulation_allcause_df, file.path(output_dir, "tables", "allcause_simulation_summary.csv"), row.names = FALSE)
write.csv(simulation_summary_df, file.path(output_dir, "tables", "decrement_simulation_summary.csv"), row.names = FALSE)
write.csv(coherence_forecast_df, file.path(output_dir, "tables", "coherence_forecast_check.csv"), row.names = FALSE)
write.csv(correlation_table, file.path(output_dir, "tables", "copula_correlation_table.csv"), row.names = FALSE)
write.csv(forecast_summary_table, file.path(output_dir, "tables", "forecast_summary_by_group.csv"), row.names = FALSE)
write.csv(share_sum_check_hist, file.path(output_dir, "tables", "historical_share_sum_check.csv"), row.names = FALSE)
write.csv(share_forecast_check, file.path(output_dir, "tables", "forecast_share_sum_check.csv"), row.names = FALSE)
if (nrow(holdout_results_df) > 0) write.csv(holdout_results_df, file.path(output_dir, "tables", "holdout_allcause_forecast.csv"), row.names = FALSE)
if (nrow(holdout_metrics_by_age) > 0) write.csv(holdout_metrics_by_age, file.path(output_dir, "tables", "holdout_allcause_metrics_by_age.csv"), row.names = FALSE)
if (nrow(holdout_metrics_overall) > 0) write.csv(holdout_metrics_overall, file.path(output_dir, "tables", "holdout_allcause_metrics_overall.csv"), row.names = FALSE)

###############################################################################
# PLOTS
###############################################################################

log_message("Writing plots...")

# 1. Historical + forecast all-cause by selected ages
hist_all_plot_df <- allcause_df %>%
  filter(age %in% selected_plot_ages) %>%
  mutate(type = "Historical", value = mx) %>%
  select(year, sex, sex_label, age, type, value)

fc_all_plot_df <- allcause_forecast_df %>%
  filter(age %in% selected_plot_ages) %>%
  mutate(type = "Forecast", value = mx_forecast_all) %>%
  select(year, sex, sex_label, age, type, value)

p_all <- bind_rows(hist_all_plot_df, fc_all_plot_df) %>%
  ggplot(aes(x = year, y = value, linetype = type, color = sex_label)) +
  geom_line() +
  facet_wrap(~ age, scales = "free_y") +
  scale_y_continuous(labels = scales::label_number(accuracy = 0.0001)) +
  labs(title = "All-cause mortality: historical and forecast", x = "Year", y = "m_x", color = "", linetype = "") +
  theme_minimal(base_size = 12)
save_plot(p_all, file.path(output_dir, "plots", "allcause_historical_forecast_v5.png"), 12, 7)

# 2. Kappa paths
p_kappa <- kappa_summary_df %>%
  ggplot(aes(x = year, y = kappa, linetype = type, color = sex_label)) +
  geom_line() +
  labs(title = "Lee-Carter kappa: historical and forecast", x = "Year", y = expression(kappa[t]), color = "", linetype = "") +
  theme_minimal(base_size = 12)
save_plot(p_kappa, file.path(output_dir, "plots", "kappa_historical_forecast.png"), 11, 6)

# 3. Forecast decrement shares at age 65
p_share <- share_forecast_df %>%
  filter(age == 65) %>%
  ggplot(aes(x = year, y = share_forecast, color = dec_group)) +
  geom_line() +
  facet_wrap(~ sex_label, scales = "free_y") +
  scale_y_continuous(labels = scales::percent_format(accuracy = 1)) +
  labs(title = "Forecast decrement shares at age 65", x = "Year", y = "Share", color = "Group") +
  theme_minimal(base_size = 12)
save_plot(p_share, file.path(output_dir, "plots", "forecast_decrement_shares_age65_v5.png"), 11, 6)

# 4. Forecast decrement rates at age 65
p_dec <- forecast_rates_df %>%
  filter(age == 65) %>%
  ggplot(aes(x = year, y = mx_forecast_dec, color = dec_group)) +
  geom_line() +
  facet_wrap(~ sex_label, scales = "free_y") +
  scale_y_continuous(labels = scales::label_number(accuracy = 0.0001)) +
  labs(title = "Forecast decrement-specific mortality at age 65", x = "Year", y = "m_x^(j)", color = "Group") +
  theme_minimal(base_size = 12)
save_plot(p_dec, file.path(output_dir, "plots", "forecast_decrement_rates_age65_v5.png"), 11, 6)

# 5. Predictive intervals for decrement mortality at age 65
if (nrow(simulation_summary_df) > 0) {
  plot_groups <- unique(simulation_summary_df$dec_group)
  plot_groups <- plot_groups[seq_len(min(4, length(plot_groups)))]
  p_int <- simulation_summary_df %>%
    filter(age == 65, dec_group %in% plot_groups) %>%
    ggplot(aes(x = year, y = median_mx)) +
    geom_ribbon(aes(ymin = low95_mx, ymax = high95_mx), alpha = 0.25) +
    geom_line() +
    facet_grid(sex_label ~ dec_group, scales = "free_y") +
    scale_y_continuous(labels = scales::label_number(accuracy = 0.0001)) +
    labs(title = "Predictive intervals for decrement mortality at age 65", x = "Year", y = "m_x^(j)") +
    theme_minimal(base_size = 12)
  save_plot(p_int, file.path(output_dir, "plots", "predictive_intervals_age65_v5.png"), 13, 8)
}

# 6. Historical age profiles for selected years
hist_profile_years <- auto_year_grid(allcause_df$year, selected_profile_years_n)
p_age_hist <- allcause_df %>%
  filter(year %in% hist_profile_years, sex %in% selected_sexes) %>%
  ggplot(aes(x = age, y = mx, color = factor(year))) +
  geom_line() +
  facet_wrap(~ sex_label, scales = "free_y") +
  scale_y_continuous(trans = "log10", labels = scales::label_number(accuracy = 0.0001)) +
  labs(title = "Historical age profiles of all-cause mortality", x = "Age", y = "m_x (log scale)", color = "Year") +
  theme_minimal(base_size = 12)
save_plot(p_age_hist, file.path(output_dir, "plots", "historical_age_profiles_allcause.png"), 11, 6)

# 7. Forecast age profiles for selected forecast years
fc_profile_years <- auto_year_grid(allcause_forecast_df$year, selected_profile_years_n)
p_age_fc <- allcause_forecast_df %>%
  filter(year %in% fc_profile_years, sex %in% selected_sexes) %>%
  ggplot(aes(x = age, y = mx_forecast_all, color = factor(year))) +
  geom_line() +
  facet_wrap(~ sex_label, scales = "free_y") +
  scale_y_continuous(trans = "log10", labels = scales::label_number(accuracy = 0.0001)) +
  labs(title = "Forecast age profiles of all-cause mortality", x = "Age", y = "Forecast m_x (log scale)", color = "Forecast year") +
  theme_minimal(base_size = 12)
save_plot(p_age_fc, file.path(output_dir, "plots", "forecast_age_profiles_allcause.png"), 11, 6)

# 8. Heatmap forecast all-cause by age/year
p_heat <- allcause_forecast_df %>%
  ggplot(aes(x = year, y = age, fill = log10(pmax(mx_forecast_all, EPS)))) +
  geom_tile() +
  facet_wrap(~ sex_label) +
  labs(title = "All-cause forecast heatmap", x = "Year", y = "Age", fill = "log10(m_x)") +
  theme_minimal(base_size = 12)
save_plot(p_heat, file.path(output_dir, "plots", "heatmap_allcause_forecast.png"), 11, 6)

# 9. Coherence check sum decrements vs all-cause at age 65
p_coh <- coherence_forecast_df %>%
  filter(age == 65) %>%
  tidyr::pivot_longer(cols = c(sum_dec_mx, allcause_mx), names_to = "series", values_to = "value") %>%
  ggplot(aes(x = year, y = value, color = series)) +
  geom_line() +
  facet_wrap(~ sex_label, scales = "free_y") +
  labs(title = "Coherence check at age 65", x = "Year", y = "Rate", color = "") +
  theme_minimal(base_size = 12)
save_plot(p_coh, file.path(output_dir, "plots", "coherence_check_age65.png"), 11, 6)

# 10. Q forecast by age for final forecast year
last_fc_year <- max(allcause_q_df$year, na.rm = TRUE)
p_q_age <- allcause_q_df %>%
  filter(year == last_fc_year, age <= plot_age_max) %>%
  ggplot(aes(x = age, y = q_forecast_all, color = sex_label)) +
  geom_line() +
  labs(title = paste("All-cause q_x forecast in", last_fc_year), x = "Age", y = "q_x", color = "") +
  theme_minimal(base_size = 12)
save_plot(p_q_age, file.path(output_dir, "plots", "allcause_q_by_age_last_forecast_year.png"), 11, 6)

# 11. Correlation heatmap of decrement shocks
if (nrow(correlation_table) > 0) {
  p_corr <- correlation_table %>%
    ggplot(aes(x = group1, y = group2, fill = correlation)) +
    geom_tile() +
    facet_wrap(~ sex_label) +
    labs(title = "Correlation matrix of decrement innovations", x = "", y = "", fill = "Corr") +
    theme_minimal(base_size = 12) +
    theme(axis.text.x = element_text(angle = 45, hjust = 1))
  save_plot(p_corr, file.path(output_dir, "plots", "copula_correlation_heatmap.png"), 12, 7)
}

# 12. Stacked share composition at age 65
p_stack <- share_forecast_df %>%
  filter(age == 65) %>%
  ggplot(aes(x = year, y = share_forecast, fill = dec_group)) +
  geom_area(position = "stack", alpha = 0.8) +
  facet_wrap(~ sex_label) +
  scale_y_continuous(labels = scales::percent_format(accuracy = 1)) +
  labs(title = "Forecast share composition at age 65", x = "Year", y = "Share", fill = "Group") +
  theme_minimal(base_size = 12)
save_plot(p_stack, file.path(output_dir, "plots", "stacked_share_composition_age65.png"), 11, 6)

# 13. All-cause forecast at selected ages up to 2070 with Males/Females labels
p_all_2070 <- bind_rows(hist_all_plot_df, fc_all_plot_df) %>%
  filter(year <= forecast_end_year) %>%
  ggplot(aes(x = year, y = value, linetype = type, color = sex_label)) +
  geom_line() +
  facet_wrap(~ age, scales = "free_y") +
  scale_y_continuous(labels = scales::label_number(accuracy = 0.0001)) +
  labs(title = paste0("All-cause mortality forecast through ", forecast_end_year), x = "Year", y = "m_x", color = "", linetype = "") +
  theme_minimal(base_size = 12)
save_plot(p_all_2070, file.path(output_dir, "plots", "allcause_historical_forecast_to_2070_v5.png"), 13, 7)

# 14. Final-year decrement composition by age
final_share_year <- max(share_forecast_df$year, na.rm = TRUE)
p_final_share_age <- share_forecast_df %>%
  filter(year == final_share_year, age <= plot_age_max) %>%
  ggplot(aes(x = age, y = share_forecast, color = dec_group)) +
  geom_line() +
  facet_wrap(~ sex_label, scales = "free_y") +
  scale_y_continuous(labels = scales::percent_format(accuracy = 1)) +
  labs(title = paste0("Decrement share profiles by age in ", final_share_year), x = "Age", y = "Share", color = "Cause group") +
  theme_minimal(base_size = 12)
save_plot(p_final_share_age, file.path(output_dir, "plots", "decrement_share_profiles_final_year_v5.png"), 12, 7)

# 15. Final-year decrement rates by age
p_final_rate_age <- forecast_rates_df %>%
  filter(year == final_share_year, age <= plot_age_max) %>%
  ggplot(aes(x = age, y = mx_forecast_dec, color = dec_group)) +
  geom_line() +
  facet_wrap(~ sex_label, scales = "free_y") +
  scale_y_continuous(trans = "log10", labels = scales::label_number(accuracy = 0.0001)) +
  labs(title = paste0("Decrement-specific mortality profiles by age in ", final_share_year), x = "Age", y = "Forecast m_x (log scale)", color = "Cause group") +
  theme_minimal(base_size = 12)
save_plot(p_final_rate_age, file.path(output_dir, "plots", "decrement_rate_profiles_final_year_v5.png"), 12, 7)

# 16. 3D surfaces for all-cause forecasts by sex
for (sx_lab in intersect(c("Males", "Females"), unique(allcause_forecast_df$sex_label))) {
  surf_df <- allcause_forecast_df %>% filter(sex_label == sx_lab, year <= forecast_end_year, age <= plot_age_max) %>%
    select(year, age, mx_forecast_all)
  save_persp_png(surf_df, x_col = "year", y_col = "age", z_col = "mx_forecast_all",
                 filename = file.path(output_dir, "plots", paste0("surface_allcause_", gsub(" ", "_", tolower(sx_lab)), "_v5.png")),
                 main_title = paste("3D surface: all-cause forecast -", sx_lab),
                 xlab = "Year", ylab = "Age", zlab = "m_x")
}

# 17. 3D surfaces for all-cause q_x by sex
for (sx_lab in intersect(c("Males", "Females"), unique(allcause_q_df$sex_label))) {
  surf_df <- allcause_q_df %>% filter(sex_label == sx_lab, year <= forecast_end_year, age <= plot_age_max) %>%
    select(year, age, q_forecast_all)
  save_persp_png(surf_df, x_col = "year", y_col = "age", z_col = "q_forecast_all",
                 filename = file.path(output_dir, "plots", paste0("surface_qx_", gsub(" ", "_", tolower(sx_lab)), "_v5.png")),
                 main_title = paste("3D surface: all-cause q_x forecast -", sx_lab),
                 xlab = "Year", ylab = "Age", zlab = "q_x")
}

# 18. 3D surface for dominant decrement share by sex
dominant_group <- share_forecast_df %>%
  group_by(dec_group) %>% summarise(avg_share = mean(share_forecast, na.rm = TRUE), .groups = "drop") %>%
  arrange(desc(avg_share)) %>% slice(1) %>% pull(dec_group)
if (length(dominant_group) == 1 && !is.na(dominant_group)) {
  for (sx_lab in intersect(c("Males", "Females"), unique(share_forecast_df$sex_label))) {
    surf_df <- share_forecast_df %>% filter(sex_label == sx_lab, dec_group == dominant_group, year <= forecast_end_year, age <= plot_age_max) %>%
      select(year, age, share_forecast)
    save_persp_png(surf_df, x_col = "year", y_col = "age", z_col = "share_forecast",
                   filename = file.path(output_dir, "plots", paste0("surface_share_", gsub(" ", "_", tolower(sx_lab)), "_v5.png")),
                   main_title = paste("3D surface: share forecast -", dominant_group, "-", sx_lab),
                   xlab = "Year", ylab = "Age", zlab = "Share")
  }
}


# 18b. Holdout plot: actual vs forecast at selected ages
if (nrow(holdout_results_df) > 0) {
  p_hold_lines <- holdout_results_df %>%
    filter(age %in% holdout_plot_ages) %>%
    tidyr::pivot_longer(cols = c(mx, mx_forecast), names_to = "series", values_to = "value") %>%
    mutate(series = dplyr::recode(series, mx = "Actual", mx_forecast = "Holdout forecast")) %>%
    ggplot(aes(x = year, y = value, color = series, linetype = series)) +
    geom_line() +
    facet_grid(sex_label ~ age, scales = "free_y") +
    scale_y_continuous(labels = scales::label_number(accuracy = 0.0001)) +
    labs(title = "Holdout validation: actual vs forecast all-cause mortality", x = "Year", y = "m_x", color = "", linetype = "") +
    theme_minimal(base_size = 12)
  save_plot(p_hold_lines, file.path(output_dir, "plots", "holdout_actual_vs_forecast_allcause.png"), 13, 7)

  p_hold_err <- holdout_error_surface_df %>%
    filter(age <= holdout_age_max) %>%
    ggplot(aes(x = year, y = age, fill = log_abs_error)) +
    geom_tile() +
    facet_wrap(~ sex_label) +
    labs(title = "Holdout validation heatmap: log10 absolute forecast error", x = "Holdout year", y = "Age", fill = "log10 abs err") +
    theme_minimal(base_size = 12)
  save_plot(p_hold_err, file.path(output_dir, "plots", "holdout_error_heatmap_allcause.png"), 11, 6)

  p_hold_mape_age <- holdout_metrics_by_age %>%
    filter(age <= holdout_age_max) %>%
    ggplot(aes(x = age, y = MAPE, color = sex_label)) +
    geom_line() +
    labs(title = "Holdout validation by age: MAPE of all-cause forecast", x = "Age", y = "MAPE", color = "") +
    theme_minimal(base_size = 12)
  save_plot(p_hold_mape_age, file.path(output_dir, "plots", "holdout_mape_by_age_allcause.png"), 11, 6)

  for (sx_lab in intersect(c("Males", "Females"), unique(holdout_results_df$sex_label))) {
    surf_act <- holdout_results_df %>% filter(sex_label == sx_lab, age <= holdout_age_max) %>% select(year, age, mx)
    save_persp_png(surf_act, x_col = "year", y_col = "age", z_col = "mx",
                   filename = file.path(output_dir, "plots", paste0("surface_holdout_actual_", gsub(" ", "_", tolower(sx_lab)), ".png")),
                   main_title = paste("3D surface: holdout actual all-cause -", sx_lab),
                   xlab = "Holdout year", ylab = "Age", zlab = "m_x")

    surf_fc <- holdout_results_df %>% filter(sex_label == sx_lab, age <= holdout_age_max) %>% select(year, age, mx_forecast)
    save_persp_png(surf_fc, x_col = "year", y_col = "age", z_col = "mx_forecast",
                   filename = file.path(output_dir, "plots", paste0("surface_holdout_forecast_", gsub(" ", "_", tolower(sx_lab)), ".png")),
                   main_title = paste("3D surface: holdout forecast all-cause -", sx_lab),
                   xlab = "Holdout year", ylab = "Age", zlab = "Forecast m_x")

    surf_err <- holdout_error_surface_df %>% filter(sex_label == sx_lab, age <= holdout_age_max) %>% select(year, age, abs_error)
    save_persp_png(surf_err, x_col = "year", y_col = "age", z_col = "abs_error",
                   filename = file.path(output_dir, "plots", paste0("surface_holdout_abs_error_", gsub(" ", "_", tolower(sx_lab)), ".png")),
                   main_title = paste("3D surface: holdout absolute error -", sx_lab),
                   xlab = "Holdout year", ylab = "Age", zlab = "Abs error")
  }
}

# 18c. Extra 3D surfaces for top decrement groups by mean share
extra_groups <- share_forecast_df %>%
  group_by(dec_group) %>% summarise(avg_share = mean(share_forecast, na.rm = TRUE), .groups = "drop") %>%
  arrange(desc(avg_share)) %>%
  filter(dec_group != "OTHER") %>%
  slice_head(n = extra_3d_groups_n) %>%
  pull(dec_group)

if (length(extra_groups) > 0) {
  for (g in extra_groups) {
    for (sx_lab in intersect(c("Males", "Females"), unique(forecast_rates_df$sex_label))) {
      surf_df_rate <- forecast_rates_df %>%
        filter(sex_label == sx_lab, dec_group == g, year <= forecast_end_year, age <= plot_age_max) %>%
        select(year, age, mx_forecast_dec)
      save_persp_png(surf_df_rate, x_col = "year", y_col = "age", z_col = "mx_forecast_dec",
                     filename = file.path(output_dir, "plots", paste0("surface_rate_", safe_label(g), "_", gsub(" ", "_", tolower(sx_lab)), ".png")),
                     main_title = paste("3D surface: decrement rate forecast -", g, "-", sx_lab),
                     xlab = "Year", ylab = "Age", zlab = "m_x^(j)")

      surf_df_q <- decrement_q_df %>%
        filter(sex_label == sx_lab, dec_group == g, year <= forecast_end_year, age <= plot_age_max) %>%
        select(year, age, q_forecast_dec)
      save_persp_png(surf_df_q, x_col = "year", y_col = "age", z_col = "q_forecast_dec",
                     filename = file.path(output_dir, "plots", paste0("surface_q_", safe_label(g), "_", gsub(" ", "_", tolower(sx_lab)), ".png")),
                     main_title = paste("3D surface: decrement q forecast -", g, "-", sx_lab),
                     xlab = "Year", ylab = "Age", zlab = "q_x^(j)")
    }
  }
}

# 19. Extra tables for 2070 snapshots
write.csv(selected_groups_table, file.path(output_dir, "tables", "selected_groups_v5.csv"), row.names = FALSE)
if (exists("cause_map")) write.csv(cause_map, file.path(output_dir, "tables", "group_definition_v5.csv"), row.names = FALSE)
write.csv(allcause_forecast_df %>% filter(year == forecast_end_year), file.path(output_dir, "tables", "allcause_forecast_2070_by_age.csv"), row.names = FALSE)
write.csv(allcause_q_df %>% filter(year == forecast_end_year), file.path(output_dir, "tables", "allcause_q_2070_by_age.csv"), row.names = FALSE)
write.csv(share_forecast_df %>% filter(year == forecast_end_year), file.path(output_dir, "tables", "share_forecast_2070_by_age.csv"), row.names = FALSE)
write.csv(forecast_rates_df %>% filter(year == forecast_end_year), file.path(output_dir, "tables", "decrement_rate_2070_by_age.csv"), row.names = FALSE)
write.csv(decrement_q_df %>% filter(year == forecast_end_year), file.path(output_dir, "tables", "decrement_q_2070_by_age.csv"), row.names = FALSE)

###############################################################################
# SAVE MODEL OBJECTS
###############################################################################

saveRDS(list(
  meta_summary = meta_summary,
  unit_diagnostics = unit_diagnostics,
  lc_results = lc_results,
  share_models = share_models,
  correlation_objects = correlation_objects,
  allcause_forecast_df = allcause_forecast_df,
  allcause_q_df = allcause_q_df,
  share_forecast_df = share_forecast_df,
  forecast_rates_df = forecast_rates_df,
  decrement_q_df = decrement_q_df,
  coherence_forecast_df = coherence_forecast_df,
  simulation_allcause_df = simulation_allcause_df,
  simulation_summary_df = simulation_summary_df,
  holdout_results_df = holdout_results_df,
  holdout_metrics_by_age = holdout_metrics_by_age,
  holdout_metrics_overall = holdout_metrics_overall
), file.path(output_dir, "rds", "dynamic_factor_copula_mortality_model_v5.rds"))

###############################################################################
# CONSOLE SUMMARY
###############################################################################

log_message("Completed successfully.")
log_message("Results workbook: %s", results_xlsx)
log_message("Plots folder: %s", file.path(output_dir, "plots"))
log_message("Tables folder: %s", file.path(output_dir, "tables"))
log_message("RDS folder: %s", file.path(output_dir, "rds"))

cat("\nDONE.\n")
cat("Source file used: ", target_file, "\n", sep = "")
cat("Output directory : ", output_dir, "\n", sep = "")
cat("Selected groups  : ", paste(selected_groups, collapse = ", "), "\n", sep = "")
cat("Forecast horizon : ", forecast_horizon, " years (through ", forecast_end_year, ")\n", sep = "")
cat("Simulations      : ", n_sims, "\n", sep = "")
cat("Unit divisor     : ", unit_scale, "\n", sep = "")
if (nrow(holdout_metrics_overall) > 0) {
  cat("Holdout years    : ", holdout_years, "\n", sep = "")
}