source("params.R")

# Imagine monthly is currently the string "True" or "true"
monthly <- toupper(monthly)        # Converts to "TRUE"
monthly <- as.logical(monthly)     # Converts to TRUE (logical)

 
slctPrdctr <- setdiff(slctPrdctr, "ALLP_INT")

options(error = quote({
  dump.frames(to.file = TRUE, dumpto = "last.dump")
  load("last.dump.rda")
  print(last.dump)
  q()
}))

###############################################################################
# ZTS Jul 2026: Added monthly inresponse to 3rd International Workshop recommendation
#ZTS MAr 2026: rewritten/optimized parallel version & Diagnostics with ChatGPT
# Keeps stepAIC inside Leave-3-out loop
# Parallelization: over grid cells
###############################################################################

library(RNetCDF)
library(nnet)
library(MASS)
library(parallel)

source("templateRPS_Wilks.R")
source("templateVerification.R")

zproctime0 <- proc.time()

###############################################################################
# DIAGNOSTIC SETTINGS
###############################################################################
RUN_DIAGNOSTIC <- TRUE 

# Number of folds to print in detail
diag_max_folds_to_print <- 15

# Write diagnostics to text file
diag_file <- paste0(workingDir,"/RPSS_diagnostic_Lead",iLead,"log.txt")

LastYrHind   <- ThisYr   #This is always true. It is model init' year
yrFCST <- ThisYr
iselLen <- 10  # tercile cats are determened for > iselLen
TmodelFileNames <- c("canesm5","gem5nemo","nasa","ccsm4","cfsv2","cesm1","ensmean")
AllmodelNamesPlot  <- c("ESM5", "GEMO","NASA","CCSM","CFS2","CESM","ENS")
names(AllmodelNamesPlot) <- TmodelFileNames

DETREND   <- FALSE
stdizeObs <- FALSE

gcmStart <- 1
gcmLast  <- length(modelFileNames)
forecastNames <- modelFileNames

modelNamesPlt <- AllmodelNamesPlot[modelFileNames]

# Parallel settings
ncores <- max(1, detectCores() - 1)
USE_MCLAPPLY <- (.Platform$OS.type == "unix")

all_nodes <- parallel::detectCores() - 1
#print(all_nodes)
# Other settings
nmiss <- 15
rthr  <- 1
xpb   <- 0.33
VARNMTHD <- paste0("_L",iLead,"_PRCP_Multinom_")
varunits <- "K"

if (monthly) {
  FTYPE  <- "_Monthly.nc"
  FTYPEF <- "_Monthly"
} else {
  # Dynamically build the string using SznLen (e.g., 3 -> "_3monSeasonal.nc")
  FTYPE  <- paste0("_", SznLen, "monSeasonal.nc")
  FTYPEF <- paste0("_", SznLen, "monSeasonal")
}

###############################################################################
# SANITY CHECKS
###############################################################################
if (length(setdiff(modelFileNames, TmodelFileNames)) != 0) stop("ModelNameError")

###############################################################################
# DIRECTORIES / FILES
###############################################################################

xdirdat  <- paste0(workingDir,"/epsfiles/")
if (!dir.exists(xdirdat)) dir.create(xdirdat)
xdirdat  <- paste0(workingDir,"/epsfreq/")
if (!dir.exists(xdirdat)) dir.create(xdirdat)

dirdat0  <- paste0(workingDir,"/analdat/")
fcstdir  <- paste0(workingDir,"/sstdat/")
hinddiro <- paste0(workingDir,"/")

if (!dir.exists(dirdat0)) dir.create(dirdat0)

hinddir  <- fcstdir

hindcastNames <- character(0)
eoftsNames    <- character(0)
eofPrecNames  <- character(0)

for (imod in gcmStart:gcmLast) {
  hindcastNames <- c(hindcastNames, paste0("sstIndex", SEASON, modelFileNames[imod], ".nc"))
  eoftsNames    <- c(eoftsNames,    paste0("eofts",   SEASON, modelFileNames[imod], ".txt"))

  eofPrecNames    <- c(eofPrecNames,    paste0("precEOFts",   SEASON, modelFileNames[imod], ".txt"))

}

month.2abb <- c(month.abb, month.abb)

###############################################################################
# PREDICTOR DEFINITIONS
###############################################################################

AllPRDCTRS <- slctPrdctr
AllhNames  <- paste0(AllPRDCTRS,"hZ")
AllfNames  <- paste0(AllPRDCTRS,"fZ")

names(AllhNames)  <- AllPRDCTRS
names(AllfNames)  <- AllPRDCTRS
names(AllPRDCTRS) <- AllPRDCTRS

PRDCTRS <- AllPRDCTRS[slctPrdctr]
prddif  <- setdiff(AllPRDCTRS, PRDCTRS)
PRDCTRS <- c(PRDCTRS, AllPRDCTRS[prddif])

hnames <- AllhNames[PRDCTRS]
fnames <- AllfNames[PRDCTRS]

nprdct  <- length(fnames)
nmodels <- length(gcmStart:gcmLast)


###############################################################################
# HELPER FUNCTIONS ;; ChatGPT rewrote my original code into functions; It shortned a great deal
#### Done by March 19, 2026.
###### Function names are:
### F1: selected_predictor_names
### F2: safe_scale_train_fcst
### F3: build_leave3_indices
### F4: obs_cat_from_binary
### F5: obs_num_from_binary
### F6: safe_predict_probs
### F7: prob_to_cat
### F8: safe_hss
### F9: row_complete
###F10: domain_diag_summary
###F11: make_predictor_matrix
###F12: safe_rps_components
###F13: cat_freq
###F14: mean_sharpness
###F15: prob_summary
###F16: extract_term_names
###F17: decode_term_names
###F18: fit_model_with_scope
###F19: run_cv_with_refit
###F20: run_cv_fixed_terms
###F21: diagnose_one_cell
###F22: cell_worker
###F23: run_serial_cells
###F24: run_parallel_cells

###############################################################################
selected_predictor_names <- function(final_model, pred_ord) {
  trms <- attr(terms(final_model), "term.labels")
  if (length(trms) == 0) return(character(0))

  idx <- suppressWarnings(as.integer(sub("^p", "", trms)))
  idx <- idx[is.finite(idx)]
  idx <- idx[idx >= 1 & idx <= length(pred_ord)]

  pred_ord[idx]
}

safe_scale_train_fcst <- function(xtrain, xfcst, do_scale = FALSE) {
  if (!do_scale) return(list(train = xtrain, fcst = xfcst))
  mu  <- mean(xtrain, na.rm = TRUE)
  sdx <- sd(xtrain, na.rm = TRUE)
  if (!is.finite(sdx) || sdx == 0) {
    return(list(train = rep(NA_real_, length(xtrain)),
                fcst  = rep(NA_real_, length(xfcst))))
  }
  list(train = (xtrain - mu) / sdx,
       fcst  = (xfcst  - mu) / sdx)
}

build_leave3_indices <- function(n) {
  out <- vector("list", n)
  for (k in seq_len(n)) {
    if (k == 1) {
      out[[k]] <- setdiff(seq_len(n), c(1, 2))
    } else if (k == n) {
      out[[k]] <- setdiff(seq_len(n), c(n - 1, n))
    } else {
      out[[k]] <- setdiff(seq_len(n), c(k - 1, k, k + 1))
    }
  }
  out
}

obs_cat_from_binary <- function(dry, normal, wet) {
  out <- rep(NA_character_, length(dry))
  out[dry == 1]    <- "dry"
  out[normal == 1] <- "normal"
  out[wet == 1]    <- "wet"
  factor(out, levels = c("dry", "normal", "wet"))
}

obs_num_from_binary <- function(dry, normal, wet) {
  dry + 2 * normal + 3 * wet
}

safe_predict_probs <- function(model, newdata) {
  pr <- predict(model, newdata = newdata, type = "probs")
  if (is.null(dim(pr))) {
    pr <- matrix(pr, nrow = 1)
  }
  colnames(pr) <- c("dry", "normal", "wet")
  pr
}

prob_to_cat0 <- function(prob_row, xpb = NULL) {
  which.max(prob_row)
}

prob_to_cat <- function(prob_input, xpb = 0.33) {
  # If it's just a vector, turn it into a 1-row matrix
  if (!is.matrix(prob_input)) {
    prob_input <- matrix(prob_input, nrow = 1)
  }

  # 1. Identify the 'winner' index for each row
  k <- max.col(prob_input, ties.method = "first")
  
  # 2. Logic: Row is invalid if:
  #    - There are any NAs
  #    - More than one value exceeds the threshold xpb
  invalid <- rowSums(is.na(prob_input)) > 0 | 
             rowSums(prob_input > xpb, na.rm = TRUE) > 1
  
  k[invalid] <- NA_integer_
  return(k)
}

prob_to_cat2 <- function(prob_row, xpb = 0.33) {
  k  <- which.max(prob_row)
  b1 <- prob_row[1]
  n1 <- prob_row[2]
  a1 <- prob_row[3]

  if (k == 1 && (a1 > xpb || n1 > xpb)) return(NA_integer_)
  if (k == 2 && (b1 > xpb || a1 > xpb)) return(NA_integer_)
  if (k == 3 && (b1 > xpb || n1 > xpb)) return(NA_integer_)
  k
}

prob_to_cat1 <- function(prob_row, xpb = 0.33) {
  k <- which_max_na(prob_row)
  if (all(prob_row[-k] <= xpb)) k else NA_integer_
}

safe_hss <- function(obs_num, probs, xpb = 0.33, obs_is_012 = FALSE) {
  if (is.null(dim(probs))) return(NA_real_)
  if (nrow(probs) < 2) return(NA_real_)

  pred_cat <- apply(probs, 1, prob_to_cat, xpb = xpb)
  #print(pred_cat)

  ok <- is.finite(obs_num) & is.finite(pred_cat)
  if (sum(ok) < 2) return(NA_real_)

  obs  <- obs_num[ok]
  if (obs_is_012) obs <- obs + 1
  pred <- pred_cat[ok]

  ok2 <- obs %in% 1:3 & pred %in% 1:3
  obs  <- obs[ok2]
  pred <- pred[ok2]

  if (length(obs) < 2) return(NA_real_)

  xx <- table(
    factor(pred, levels = 1:3),
    factor(obs,  levels = 1:3)
  )

  yy <- rep(1/3, 3)
  Acat <- multi.cont(xx, baseline = yy)
  
  Acat$hss
}

row_complete <- function(x) {
  rowSums(is.finite(x)) == ncol(x)
}
safe_mean <- function(x) {
  x <- x[is.finite(x)]
  if (length(x) == 0) NA_real_ else mean(x)
}

safe_median <- function(x) {
  x <- x[is.finite(x)]
  if (length(x) == 0) NA_real_ else median(x)
}

safe_max <- function(x) {
  x <- x[is.finite(x)]
  if (length(x) == 0) NA_real_ else max(x)
}

safe_frac <- function(cond) {
  cond <- cond[!is.na(cond)]
  if (length(cond) == 0) NA_real_ else mean(cond)
}

domain_diag_summary <- function(ans, label = "") {
  cat("\n================ DOMAIN SUMMARY", label, "================\n")

cat("NA hhss count     :", sum(is.na(ans$hhss)), "\n")
cat("Finite hhss count :", sum(is.finite(ans$hhss)), "\n")

  cat("Valid hrpss cells :", sum(is.finite(ans$hrpss)), "\n")
  cat("Valid hhss cells  :", sum(is.finite(ans$hhss)), "\n")
  cat("Valid n_selected  :", sum(is.finite(ans$n_selected)), "\n")

  cat("Mean hrpss        :", safe_mean(ans$hrpss), "\n")
  cat("Median hrpss      :", safe_median(ans$hrpss), "\n")
  cat("Max hrpss         :", safe_max(ans$hrpss), "\n")
  cat("Frac hrpss < 0    :", safe_frac(ans$hrpss < 0), "\n")
  cat("Frac hrpss < .05  :", safe_frac(ans$hrpss < 0.05), "\n")
  cat("Frac hrpss >= .10 :", safe_frac(ans$hrpss >= 0.10), "\n")
  cat("Frac hrpss >= .20 :", safe_frac(ans$hrpss >= 0.20), "\n")

  cat("Mean hhss         :", safe_mean(ans$hhss), "\n")
  cat("Mean n_selected   :", safe_mean(ans$n_selected), "\n")
}
####### make_predictor_matrix
make_predictor_matrix <- function(hind_list, fcst_list, pred_order, imodel) {
  nh <- nrow(hind_list[[pred_order[1]]])
  nf <- nrow(fcst_list[[pred_order[1]]])

  Xh <- matrix(NA_real_, nrow = nh, ncol = length(pred_order))
  Xf <- matrix(NA_real_, nrow = nf, ncol = length(pred_order))

  for (k in seq_along(pred_order)) {
    pn <- pred_order[k]

    if (is.null(hind_list[[pn]])) stop("Missing hind_list entry for predictor: ", pn)
    if (is.null(fcst_list[[pn]])) stop("Missing fcst_list entry for predictor: ", pn)

    xh <- hind_list[[pn]][, imodel]
    xf <- fcst_list[[pn]][, imodel]

    if (length(xh) != nh) {
      stop("Training length mismatch for predictor ", pn,
           ": length(xh)=", length(xh), " but nh=", nh)
    }

    if (length(xf) != nf) {
      stop("Forecast length mismatch for predictor ", pn,
           ": length(xf)=", length(xf), " but nf=", nf)
    }

    ispc <- grepl("^PC", pn)
    z <- safe_scale_train_fcst(xh, xf, do_scale = ispc)

    if (length(z$train) != nh) {
      stop("Scaled training length mismatch for predictor ", pn,
           ": length(z$train)=", length(z$train), " but nh=", nh)
    }

    if (length(z$fcst) != nf) {
      stop("Scaled forecast length mismatch for predictor ", pn,
           ": length(z$fcst)=", length(z$fcst), " but nf=", nf)
    }

    Xh[, k] <- z$train
    Xf[, k] <- z$fcst
  }

  colnames(Xh) <- paste0("p", seq_along(pred_order))
  colnames(Xf) <- paste0("p", seq_along(pred_order))

  list(Xh = Xh, Xf = Xf)
}

######### make_predictor_matrix
make_predictor_matrix2 <- function(hind_list, fcst_list, pred_order, imodel) {
  nh <- nrow(hind_list[[pred_order[1]]])
  nf <- nrow(fcst_list[[pred_order[1]]])

  Xh <- matrix(NA_real_, nrow = nh, ncol = length(pred_order))
  Xf <- matrix(NA_real_, nrow = nf, ncol = length(pred_order))

  for (k in seq_along(pred_order)) {
    pn <- pred_order[k]
    xh <- hind_list[[pn]][, imodel]
    xf <- fcst_list[[pn]][, imodel]

    ispc <- grepl("^PC", pn)
    z <- safe_scale_train_fcst(xh, xf, do_scale = ispc)

    Xh[, k] <- z$train
    Xf[, k] <- z$fcst
  }

  colnames(Xh) <- paste0("p", seq_along(pred_order))
  colnames(Xf) <- paste0("p", seq_along(pred_order))

  list(Xh = Xh, Xf = Xf)
}

###############################################################################
# DIAGNOSTIC HELPERS
###############################################################################
safe_rps_components <- function(obs_num, probs) {
  # obs_num in {1,2,3}; probs is n x 3
  ok <- is.finite(obs_num) & rowSums(is.finite(probs)) == 3
  obs_num <- obs_num[ok]
  probs   <- probs[ok, , drop = FALSE]

  if (length(obs_num) < 2) {
    return(list(
      n = length(obs_num),
      rps_mean = NA_real_,
      clim_rps_mean = NA_real_,
      rpss = NA_real_,
      rps_each = NA_real_,
      clim_rps_each = NA_real_
    ))
  }

  # Observed cumulative vectors
  O <- matrix(0, nrow = length(obs_num), ncol = 3)
  for (i in seq_along(obs_num)) {
    O[i, obs_num[i]] <- 1
  }
  Ocum <- t(apply(O, 1, cumsum))

  Pcum <- t(apply(probs, 1, cumsum))

  # Climatology from observed frequencies
  freq <- tabulate(obs_num, nbins = 3) / length(obs_num)
  Ccum <- matrix(rep(cumsum(freq), each = length(obs_num)), ncol = 3)

  # Standard 3-category RPS uses first K-1 cumulative bins
  rps_each      <- rowSums((Pcum[, 1:2, drop = FALSE] - Ocum[, 1:2, drop = FALSE])^2)
  clim_rps_each <- rowSums((Ccum[, 1:2, drop = FALSE] - Ocum[, 1:2, drop = FALSE])^2)

  rps_mean      <- mean(rps_each)
  clim_rps_mean <- mean(clim_rps_each)
  rpss <- 1 - (rps_mean / clim_rps_mean)

  list(
    n = length(obs_num),
    rps_mean = rps_mean,
    clim_rps_mean = clim_rps_mean,
    rpss = rpss,
    rps_each = rps_each,
    clim_rps_each = clim_rps_each
  )
}

cat_freq <- function(obs_num) {
  ok <- is.finite(obs_num)
  if (sum(ok) == 0) return(c(dry = NA, normal = NA, wet = NA))
  p <- tabulate(obs_num[ok], nbins = 3) / sum(ok)
  names(p) <- c("dry", "normal", "wet")
  p
}

mean_sharpness <- function(probs) {
  if (is.null(dim(probs))) return(NA_real_)
  ok <- rowSums(is.finite(probs)) == 3
  if (sum(ok) == 0) return(NA_real_)
  mean(apply(probs[ok, , drop = FALSE], 1, max))
}

compute_sharpness <- function(probs) {
  if (is.null(dim(probs))) return(NA_real_)
  if (nrow(probs) < 1) return(NA_real_)

  mean(apply(probs, 1, max), na.rm = TRUE)
}

prob_summary <- function(probs) {
  if (is.null(dim(probs))) {
    return(c(mean_dry = NA, mean_normal = NA, mean_wet = NA, sharpness = NA))
  }
  ok <- rowSums(is.finite(probs)) == 3
  if (sum(ok) == 0) {
    return(c(mean_dry = NA, mean_normal = NA, mean_wet = NA, sharpness = NA))
  }
  c(
    mean_dry    = mean(probs[ok, 1]),
    mean_normal = mean(probs[ok, 2]),
    mean_wet    = mean(probs[ok, 3]),
    sharpness   = mean_sharpness(probs[ok, , drop = FALSE])
  )
}

extract_term_names <- function(model) {
  trms <- attr(terms(model), "term.labels")
  if (length(trms) == 0) return(character(0))
  trms
}

decode_term_names <- function(term_names, pred_ord) {
  if (length(term_names) == 0) return(character(0))
  idx <- suppressWarnings(as.integer(sub("^p", "", term_names)))
  idx <- idx[is.finite(idx)]
  idx <- idx[idx >= 1 & idx <= length(pred_ord)]
  pred_ord[idx]
}

fit_model_with_scope <- function(dataH, pred_ord, lower_type = c("forced_p1", "intercept_only")) {
  lower_type <- match.arg(lower_type)

  pcols <- paste0("p", seq_len(ncol(dataH) - 1))
  if (lower_type == "forced_p1") {
    lower_formula <- as.formula("category ~ p1")
  } else {
    lower_formula <- as.formula("category ~ 1")
  }
  upper_formula <- as.formula(paste("category ~", paste(pcols, collapse = " + ")))

  simlMod <- try(multinom(lower_formula, data = dataH, trace = FALSE), silent = TRUE)
  if (inherits(simlMod, "try-error")) return(NULL)

  final_model <- try(
    stepAIC(simlMod,
            scope = list(lower = lower_formula, upper = upper_formula),
            direction = "forward",
            trace = 0),
    silent = TRUE
  )
  if (inherits(final_model, "try-error")) return(NULL)

  list(
    model = final_model,
    lower_formula = lower_formula,
    upper_formula = upper_formula,
    selected_terms = extract_term_names(final_model),
    selected_names = decode_term_names(extract_term_names(final_model), pred_ord)
  )
} # fit_model_with_scope

run_cv_with_refit <- function(dataH, leave_idx, lower_type = c("forced_p1", "intercept_only"),
                              pred_ord, max_print = 10) {
  lower_type <- match.arg(lower_type)
  ncv <- nrow(dataH)
  loopobs <- matrix(NA_real_, nrow = ncv, ncol = 4)
  colnames(loopobs) <- c("obs", "dry", "normal", "wet")

  selected_each_fold <- vector("list", ncv)

  pcols <- paste0("p", seq_len(ncol(dataH) - 1))
  if (lower_type == "forced_p1") {
    lower_formula <- as.formula("category ~ p1")
  } else {
    lower_formula <- as.formula("category ~ 1")
  }
  upper_formula <- as.formula(paste("category ~", paste(pcols, collapse = " + ")))

  for (iloo in seq_len(ncv)) {
    isel <- leave_idx[[iloo]]
    isel <- isel[isel <= ncv]
    if (length(isel) < iselLen) next

    dtrain <- dataH[isel, , drop = FALSE]
    dtest  <- dataH[iloo, , drop = FALSE]

    simlMod <- try(multinom(lower_formula, data = dtrain, trace = FALSE), silent = TRUE)
    if (inherits(simlMod, "try-error")) next

    final_model <- try(
      stepAIC(simlMod,
              scope = list(lower = lower_formula, upper = upper_formula),
              direction = "forward",
              trace = 0),
      silent = TRUE
    )
    if (inherits(final_model, "try-error")) next

    lfcst <- try(safe_predict_probs(final_model, dtest), silent = TRUE)
    if (inherits(lfcst, "try-error")) next

    selected_each_fold[[iloo]] <- decode_term_names(extract_term_names(final_model), pred_ord)
    loopobs[iloo, 1]   <- as.numeric(dtest$category)
    loopobs[iloo, 2:4] <- lfcst[1, ]
  }

  ok <- row_complete(loopobs)
  out <- loopobs[ok, , drop = FALSE]

  list(
    loopobs = out,
    selected_each_fold = selected_each_fold,
    printed_folds = min(ncv, max_print)
  )
} # run_cv_with_refit

run_cv_fixed_terms <- function(dataH, leave_idx, fixed_terms) {
  ncv <- nrow(dataH)
  loopobs <- matrix(NA_real_, nrow = ncv, ncol = 4)
  colnames(loopobs) <- c("obs", "dry", "normal", "wet")

  if (length(fixed_terms) == 0) {
    fit_formula <- as.formula("category ~ 1")
  } else {
    fit_formula <- as.formula(paste("category ~", paste(fixed_terms, collapse = " + ")))
  }

  for (iloo in seq_len(ncv)) {
    isel <- leave_idx[[iloo]]
    isel <- isel[isel <= ncv]
    if (length(isel) < iselLen) next

    dtrain <- dataH[isel, , drop = FALSE]
    dtest  <- dataH[iloo, , drop = FALSE]

    fit <- try(multinom(fit_formula, data = dtrain, trace = FALSE), silent = TRUE)
    if (inherits(fit, "try-error")) next

    pr <- try(safe_predict_probs(fit, dtest), silent = TRUE)
    if (inherits(pr, "try-error")) next

    loopobs[iloo, 1]   <- as.numeric(dtest$category)
    loopobs[iloo, 2:4] <- pr[1, ]
  }

  ok <- row_complete(loopobs)
  loopobs[ok, , drop = FALSE]
} # run_cv_fixed_terms

#### ALL Reliability functions
reliability_one_category <- function(obs_num, probs, category = 1,
                                     breaks = seq(0, 1, by = 0.1),
                                     min_bin_count = 1) {
  # obs_num : integer vector in {1,2,3}
  # probs   : matrix/data.frame with 3 columns of forecast probs
  # category: 1, 2, or 3
  # breaks  : probability bin edges
  # returns : reliability table + BS/BSS + decomposition

  probs <- as.matrix(probs)

  ok <- is.finite(obs_num) &
        obs_num %in% 1:3 &
        rowSums(is.finite(probs)) == 3

  obs_num <- obs_num[ok]
  probs   <- probs[ok, , drop = FALSE]

  if (length(obs_num) == 0) {
    return(list(
      category = category,
      n = 0,
      clim = NA_real_,
      bs = NA_real_,
      bs_clim = NA_real_,
      bss = NA_real_,
      reliability = NA_real_,
      resolution = NA_real_,
      uncertainty = NA_real_,
      table = data.frame()
    ))
  }

  p <- probs[, category]
  y <- as.integer(obs_num == category)

  # Climatology of event
  clim <- mean(y)

  # Brier score
  bs <- mean((p - y)^2)

  # Climatological Brier score
  bs_clim <- mean((clim - y)^2)

  # Brier skill score
  bss <- if (isTRUE(all.equal(bs_clim, 0))) NA_real_ else 1 - bs / bs_clim

  # Assign bins
  bin_id <- cut(p, breaks = breaks, include.lowest = TRUE, right = TRUE, labels = FALSE)

  # Reliability table
  nb <- length(breaks) - 1
  rel_tab <- data.frame(
    bin = seq_len(nb),
    p_lo = breaks[-length(breaks)],
    p_hi = breaks[-1],
    n_bin = 0,
    p_mean = NA_real_,
    obs_freq = NA_real_,
    stringsAsFactors = FALSE
  )

  for (b in seq_len(nb)) {
    ii <- which(bin_id == b)
    rel_tab$n_bin[b] <- length(ii)

    if (length(ii) >= min_bin_count) {
      rel_tab$p_mean[b]   <- mean(p[ii])
      rel_tab$obs_freq[b] <- mean(y[ii])
    }
  }

  # Murphy decomposition
  # BS = reliability - resolution + uncertainty
  N <- length(y)
  ybar <- clim

  reliability <- 0
  resolution  <- 0

  for (b in seq_len(nb)) {
    nb_i <- rel_tab$n_bin[b]
    if (nb_i >= min_bin_count) {
      pk <- rel_tab$p_mean[b]
      okk <- rel_tab$obs_freq[b]
      wk <- nb_i / N

      reliability <- reliability + wk * (pk - okk)^2
      resolution  <- resolution  + wk * (okk - ybar)^2
    }
  }

  uncertainty <- ybar * (1 - ybar)

  list(
    category = category,
    n = N,
    clim = clim,
    bs = bs,
    bs_clim = bs_clim,
    bss = bss,
    reliability = reliability,
    resolution = resolution,
    uncertainty = uncertainty,
    table = rel_tab
  )
} # reliability_one_category

##reliability_tercile

reliability_tercile <- function(obs_num, probs,
                                breaks = seq(0, 1, by = 0.1),
                                min_bin_count = 1,
                                category_names = c("below", "normal", "above")) {
  out <- vector("list", 3)

  for (k in 1:3) {
    rr <- reliability_one_category(
      obs_num = obs_num,
      probs = probs,
      category = k,
      breaks = breaks,
      min_bin_count = min_bin_count
    )
    rr$category_name <- category_names[k]
    out[[k]] <- rr
  }

  names(out) <- category_names
  out
}

plot_reliability_diagram <- function(rel_one,
                                     main = NULL,
                                     xlab = "Forecast probability",
                                     ylab = "Observed relative frequency") {
  tab <- rel_one$table
  ok <- is.finite(tab$p_mean) & is.finite(tab$obs_freq) & tab$n_bin > 0

  if (!any(ok)) {
    plot.new()
    title(main = if (is.null(main)) paste("Reliability:", rel_one$category_name) else main)
    text(0.5, 0.5, "No valid bins")
    return(invisible(NULL))
  }

  plot(tab$p_mean[ok], tab$obs_freq[ok],
       xlim = c(0, 1), ylim = c(0, 1),
       xlab = xlab, ylab = ylab,
       main = if (is.null(main)) paste("Reliability:", rel_one$category_name) else main,
       pch = 16)
  abline(0, 1, lty = 2)
  text(tab$p_mean[ok], tab$obs_freq[ok],
       labels = tab$n_bin[ok], pos = 3, cex = 0.8)
} # reliability_tercile

############
diagnose_one_cell <- function(ilon, jlat,
                              Xh, Xf, pred_ord,
                              oblw, onrm, oabv, tblw, tnrm, tabv,
                              leave_idx,
                              trainYrsObs, trainYrsMod, testYrsObs, testYrsMod, fcstyrsObs, fcstyrsMod,
                              xpb, nmiss,
                              max_print = 10,
                              con = stdout()) {

  d_oblw <- dim(oblw)
  d_onrm <- dim(onrm)
  d_oabv <- dim(oabv)
  d_tblw <- dim(tblw)
  d_tnrm <- dim(tnrm)
  d_tabv <- dim(tabv)

  cat("\n============================================================\n", file = con)
  cat(sprintf("DIAGNOSTIC FOR GRID CELL: ilon=%d, jlat=%d\n", ilon, jlat), file = con)
  cat("============================================================\n", file = con)
  cat("dim(oblw) = ", paste(d_oblw, collapse = " x "), "\n", file = con)
  cat("dim(tblw) = ", paste(d_tblw, collapse = " x "), "\n", file = con)

  # Basic checks
  if (length(d_oblw) != 3 || length(d_tblw) != 3) {
    cat("ERROR: Expected 3-D arrays [lon, lat, time].\n", file = con)
    return(invisible(NULL))
  }

  if (ilon < 1 || ilon > d_oblw[1] || jlat < 1 || jlat > d_oblw[2]) {
    cat(sprintf("ERROR: Requested (ilon=%d, jlat=%d) is outside obs dims [1:%d, 1:%d]\n",
                ilon, jlat, d_oblw[1], d_oblw[2]), file = con)
    return(invisible(NULL))
  }

  if (ilon < 1 || ilon > d_tblw[1] || jlat < 1 || jlat > d_tblw[2]) {
    cat(sprintf("ERROR: Requested (ilon=%d, jlat=%d) is outside test dims [1:%d, 1:%d]\n",
                ilon, jlat, d_tblw[1], d_tblw[2]), file = con)
    return(invisible(NULL))
  }

  dry  <- oblw[ilon, jlat, ]
  nrm  <- onrm[ilon, jlat, ]
  wet  <- oabv[ilon, jlat, ]

  tdry <- tblw[ilon, jlat, ]
  tavg <- tnrm[ilon, jlat, ]
  twet <- tabv[ilon, jlat, ]

  obs_cat <- obs_cat_from_binary(dry, nrm, wet)
  obs_num <- obs_num_from_binary(dry, nrm, wet)

  ok <- is.finite(obs_num) & row_complete(Xh)

  cat("\n--- YEAR ALIGNMENT ---\n", file = con)
  cat("trainYrsObs: ", paste(trainYrsObs, collapse = ","), "\n", file = con)
  cat("trainYrsMod: ", paste(trainYrsMod, collapse = ","), "\n", file = con)
  cat("testYrsObs : ", paste(testYrsObs, collapse = ","), "\n", file = con)
  cat("testYrsMod : ", paste(testYrsMod, collapse = ","), "\n", file = con)
  cat("fcstyrsObs : ", paste(fcstyrsObs, collapse = ","), "\n", file = con)
  cat("fcstyrsMod : ", paste(fcstyrsMod, collapse = ","), "\n", file = con)

  cat("\n--- SAMPLE COUNTS ---\n", file = con)
  cat("training length total      :", length(obs_num), "\n", file = con)
  cat("training length valid (ok) :", sum(ok), "\n", file = con)
  cat("minimum required nmiss     :", nmiss, "\n", file = con)

  if (sum(ok) < nmiss) {
    cat("Not enough valid data for diagnosis.\n", file = con)
    return(invisible(NULL))
  }

  obs_cat_ok <- obs_cat[ok]
  obs_num_ok <- obs_num[ok]
  Xh_ok      <- Xh[ok, , drop = FALSE]

  dataH <- data.frame(category = obs_cat_ok, Xh_ok, check.names = FALSE)
  dataF <- data.frame(Xf, check.names = FALSE)

  cat("\n--- OBSERVED TRAINING CATEGORY FREQUENCIES ---\n", file = con)
  print(round(cat_freq(obs_num_ok), 4), file = con)

  # 1) forced p1 model
  fit_forced <- fit_model_with_scope(dataH, pred_ord, lower_type = "forced_p1")
  # 2) intercept-only start
  fit_int <- fit_model_with_scope(dataH, pred_ord, lower_type = "intercept_only")

  cat("\n--- FULL-TRAIN SELECTED PREDICTORS ---\n", file = con)
  cat("Forced p1 start     :", if (is.null(fit_forced)) "FIT FAILED" else paste(fit_forced$selected_names, collapse = ","), "\n", file = con)
  cat("Intercept-only start:", if (is.null(fit_int)) "FIT FAILED" else paste(fit_int$selected_names, collapse = ","), "\n", file = con)

  # Cross-validation with re-selection every fold
  cv_forced <- run_cv_with_refit(dataH, leave_idx, lower_type = "forced_p1",
                                 pred_ord = pred_ord, max_print = max_print)
  cv_int <- run_cv_with_refit(dataH, leave_idx, lower_type = "intercept_only",
                              pred_ord = pred_ord, max_print = max_print)

  cat("\n--- FOLD-BY-FOLD SELECTED PREDICTORS (first folds) ---\n", file = con)
  for (k in seq_len(min(length(cv_forced$selected_each_fold), max_print))) {
    s1 <- cv_forced$selected_each_fold[[k]]
    s2 <- cv_int$selected_each_fold[[k]]
    cat(sprintf("fold %02d | forced_p1: %s | intercept_only: %s\n",
                k,
                if (length(s1) == 0) "" else paste(s1, collapse = ","),
                if (length(s2) == 0) "" else paste(s2, collapse = ",")),
        file = con)
  }

  cat("\n--- CV PROBABILITY SUMMARY ---\n", file = con)
  if (nrow(cv_forced$loopobs) > 0) {
    print(round(prob_summary(cv_forced$loopobs[, 2:4, drop = FALSE]), 4), file = con)
  } else {
    cat("forced_p1 CV failed\n", file = con)
  }
  if (nrow(cv_int$loopobs) > 0) {
    print(round(prob_summary(cv_int$loopobs[, 2:4, drop = FALSE]), 4), file = con)
  } else {
    cat("intercept_only CV failed\n", file = con)
  }

  cat("\n--- CV RPSS DECOMPOSITION ---\n", file = con)
  if (nrow(cv_forced$loopobs) > 3) {
    dec1 <- safe_rps_components(cv_forced$loopobs[, 1], cv_forced$loopobs[, 2:4, drop = FALSE])
    cat("forced_p1 refit-each-fold:\n", file = con)
    print(round(unlist(dec1[c("n","rps_mean","clim_rps_mean","rpss")]), 5), file = con)

    rel1 <- reliability_tercile(
        obs_num = cv_forced$loopobs[, 1],
        probs   = cv_forced$loopobs[, 2:4, drop = FALSE],
        min_bin_count = 2
    )

    cat("\n--- RELIABILITY SUMMARY: forced_p1 CV ---\n", file = con)
    for (nm in names(rel1)) {
      rr <- rel1[[nm]]
      cat(nm, ": BSS=", round(rr$bss, 4),
        " Rel=", round(rr$reliability, 4),
        " Res=", round(rr$resolution, 4),
        " Unc=", round(rr$uncertainty, 4), "\n",
        file = con)
    }

  } else {
    cat("forced_p1 refit-each-fold: insufficient valid CV rows\n", file = con)
  }

  if (nrow(cv_int$loopobs) > 3) {
    dec2 <- safe_rps_components(cv_int$loopobs[, 1], cv_int$loopobs[, 2:4, drop = FALSE])
    cat("intercept_only refit-each-fold:\n", file = con)
    print(round(unlist(dec2[c("n","rps_mean","clim_rps_mean","rpss")]), 5), file = con)

    rel1 <- reliability_tercile(
        obs_num = cv_forced$loopobs[, 1],
        probs   = cv_forced$loopobs[, 2:4, drop = FALSE],
        min_bin_count = 2
    )

    cat("\n--- RELIABILITY SUMMARY: intercept_only CV ---\n", file = con)
    for (nm in names(rel1)) {
      rr <- rel1[[nm]]
      cat(nm, ": BSS=", round(rr$bss, 4),
        " Rel=", round(rr$reliability, 4),
        " Res=", round(rr$resolution, 4),
        " Unc=", round(rr$uncertainty, 4), "\n",
        file = con)
    }

  } else {
    cat("intercept_only refit-each-fold: insufficient valid CV rows\n", file = con)
  }

  # Fixed-term CV using terms chosen from full training fit
  cat("\n--- FIXED-TERM CV TEST ---\n", file = con)
  if (!is.null(fit_forced)) {
    fixed_cv1 <- run_cv_fixed_terms(dataH, leave_idx, fit_forced$selected_terms)
    if (nrow(fixed_cv1) > 3) {
      dec3 <- safe_rps_components(fixed_cv1[, 1], fixed_cv1[, 2:4, drop = FALSE])
      cat("forced_p1 selected terms, fixed across folds:\n", file = con)
      print(round(unlist(dec3[c("n","rps_mean","clim_rps_mean","rpss")]), 5), file = con)
      print(round(prob_summary(fixed_cv1[, 2:4, drop = FALSE]), 4), file = con)
    } else {
      cat("forced_p1 fixed-term CV failed or insufficient rows\n", file = con)
    }
  }

  if (!is.null(fit_int)) {
    fixed_cv2 <- run_cv_fixed_terms(dataH, leave_idx, fit_int$selected_terms)
    if (nrow(fixed_cv2) > 3) {
      dec4 <- safe_rps_components(fixed_cv2[, 1], fixed_cv2[, 2:4, drop = FALSE])
      cat("intercept_only selected terms, fixed across folds:\n", file = con)
      print(round(unlist(dec4[c("n","rps_mean","clim_rps_mean","rpss")]), 5), file = con)
      print(round(prob_summary(fixed_cv2[, 2:4, drop = FALSE]), 4), file = con)
    } else {
      cat("intercept_only fixed-term CV failed or insufficient rows\n", file = con)
    }
  }

  # Full-training forecast distribution over test/forecast rows
  cat("\n--- FULL-FIT FORECAST SUMMARY OVER dataF ---\n", file = con)
  if (!is.null(fit_forced)) {
    prf <- try(safe_predict_probs(fit_forced$model, dataF), silent = TRUE)
    if (!inherits(prf, "try-error")) {
      cat("forced_p1:\n", file = con)
      print(round(prob_summary(prf), 4), file = con)
    }
  }
  if (!is.null(fit_int)) {
    pri <- try(safe_predict_probs(fit_int$model, dataF), silent = TRUE)
    if (!inherits(pri, "try-error")) {
      cat("intercept_only:\n", file = con)
      print(round(prob_summary(pri), 4), file = con)
    }
  }

  # Full-training forecast distribution over forecast rows
  cat("\n--- FULL-FIT FORECAST SUMMARY OVER dataF ---\n", file = con)
  if (!is.null(fit_forced)) {
    prf <- try(safe_predict_probs(fit_forced$model, dataF), silent = TRUE)
    if (!inherits(prf, "try-error")) {
      cat("forced_p1:\n", file = con)
      print(round(prob_summary(prf), 4), file = con)
    }
  }
  if (!is.null(fit_int)) {
    pri <- try(safe_predict_probs(fit_int$model, dataF), silent = TRUE)
    if (!inherits(pri, "try-error")) {
      cat("intercept_only:\n", file = con)
      print(round(prob_summary(pri), 4), file = con)
    }
  }

  cat("\n--- QUICK INTERPRETATION RULES (HINDCAST CV / hRPSS) ---\n", file = con)
  cat("* If mean_dry/mean_normal/mean_wet stay near observed climatology and sharpness is near 0.33-0.40, CV forecasts are not sharp.\n", file = con)
  cat("* If intercept_only CV RPSS >> forced_p1 CV RPSS, forcing p1 is hurting hRPSS.\n", file = con)
  cat("* If fixed-term CV RPSS >> refit-each-fold CV RPSS, fold-by-fold stepAIC instability is hurting hRPSS.\n", file = con)
  cat("* If trainYrsObs and trainYrsMod do not line up as intended, year alignment may be hurting CV skill.\n", file = con)
  cat("* If category frequencies are very imbalanced or valid n is small, hRPSS will be noisy and often low.\n", file = con)

  invisible(NULL)
} # diagnose_one_cell

####################################
###### Cell Worker
######################################
cell_worker <- function(cell_id,
                        nlon, nlat,
                        oblw, onrm, oabv,
                        tblw, tnrm, tabv,
                        zblw, znrm, zabv,
                        Xh, Xf,zXh, zXf,
                        nmiss, xpb, Leave3Out,
                        leave_idx,
                        pred_ord,
                        RUN_TYPE) {

  ilon <- ((cell_id - 1) %% nlon) + 1
  jlat <- ((cell_id - 1) %/% nlon) + 1

  dry <- oblw[ilon, jlat, ]
  nrm <- onrm[ilon, jlat, ]
  wet <- oabv[ilon, jlat, ]

  obs_cat <- obs_cat_from_binary(dry, nrm, wet)
  obs_num <- obs_num_from_binary(dry, nrm, wet)

  tdry <- tblw[ilon, jlat, ]
  tavg <- tnrm[ilon, jlat, ]
  twet <- tabv[ilon, jlat, ]

  xdry <- zblw[ilon, jlat, ]
  xnrm <- znrm[ilon, jlat, ]
  xwet <- zabv[ilon, jlat, ]

  zobs_cat <- obs_cat_from_binary(xdry, xnrm, xwet)
  zobs_num <- obs_num_from_binary(xdry, xnrm, xwet)

  selected_str <- NA_character_
  n_selected   <- NA_real_

  test_probs_vec <- rep(NA_real_, 3 * length(tdry))

  ok <- is.finite(obs_num) & row_complete(Xh)
  zok <- is.finite(zobs_num) & row_complete(zXh)

  if (sum(ok) < nmiss) {
    return(list(
      scores = c(below = NA_real_, normal = NA_real_, above = NA_real_,
                 hrpss = NA_real_, trpss = NA_real_, hhss = NA_real_,
		 hsharpness = NA_real_),
      n_selected = n_selected,
      selected_str = selected_str
    ))
  }

  obs_cat_ok <- obs_cat[ok]
  obs_num_ok <- obs_num[ok]
  Xh_ok      <- Xh[ok, , drop = FALSE]

  zobs_cat_ok <- zobs_cat[zok]
  zobs_num_ok <- zobs_num[zok]
  zXh_ok      <- zXh[zok, , drop = FALSE]

  if (nrow(Xh_ok) <= 10) {
    return(list(
      scores = c(below = NA_real_, normal = NA_real_, above = NA_real_,
                 hrpss = NA_real_, trpss = NA_real_, hhss = NA_real_,
		 hsharpness = NA_real_),
      n_selected = n_selected,
      selected_str = selected_str
    ))
  }

  dataH <- data.frame(category = obs_cat_ok, Xh_ok, check.names = FALSE)
  dataF <- data.frame(Xf, check.names = FALSE)

  dataHz <- data.frame(category = zobs_cat_ok, zXh_ok, check.names = FALSE)
  dataFz <- data.frame(zXf, check.names = FALSE)

  pcols <- paste0("p", seq_len(ncol(Xh_ok)))
  upper_formula <- as.formula(paste("category ~", paste(pcols, collapse = " + ")))

  if (RUN_TYPE == "forced_p1") {
    lower_formula <- as.formula("category ~ p1")
  } else if (RUN_TYPE == "intercept_all") {
    lower_formula <- as.formula("category ~ 1")
  } else {
    stop("Unknown RUN_TYPE: ", RUN_TYPE)
  }

  hrpss <- NA_real_
  trpss <- NA_real_
  hhss  <- NA_real_
  below <- NA_real_
  normal <- NA_real_
  above <- NA_real_
  hsharpness <- NA_real_

  if (Leave3Out) {
    ncv <- nrow(dataH)
    loopobs <- matrix(NA_real_, nrow = ncv, ncol = 4)
    colnames(loopobs) <- c("obs", "dry", "normal", "wet")

    for (iloo in seq_len(ncv)) {
      isel <- leave_idx[[iloo]]
      isel <- isel[isel <= ncv]

      if (length(isel) < iselLen) next

      dtrain <- dataH[isel, , drop = FALSE]
      dtest  <- dataH[iloo, , drop = FALSE]

      simlMod <- try(multinom(lower_formula, data = dtrain, trace = FALSE), silent = TRUE)
      if (inherits(simlMod, "try-error")) next

      final_model <- try(
        stepAIC(simlMod,
                scope = list(lower = lower_formula, upper = upper_formula),
                direction = "forward",
                trace = 0),
        silent = TRUE
      )
      if (inherits(final_model, "try-error")) next

      lfcst <- try(safe_predict_probs(final_model, dtest), silent = TRUE)
      if (inherits(lfcst, "try-error")) next

      loopobs[iloo, 1]   <- obs_num_ok[iloo]
      loopobs[iloo, 2:4] <- lfcst[1, ]
    }

    okcv <- row_complete(loopobs)
    loopobs2 <- loopobs[okcv, , drop = FALSE]

    if (nrow(loopobs2) > 3) {
      hrpss <- zrpss(loopobs2)
      hhss  <- safe_hss(loopobs2[, 1], loopobs2[, 2:4, drop = FALSE], xpb = xpb)
      tmp <- loopobs2[,2:4]
      hsharpness <- mean_sharpness(tmp) 

    }

    simlMod <- try(multinom(lower_formula, data = dataH, trace = FALSE), silent = TRUE)
    if (!inherits(simlMod, "try-error")) {
      final_model <- try(
        stepAIC(simlMod,
                scope = list(lower = lower_formula, upper = upper_formula),
                direction = "forward",
                trace = 0),
        silent = TRUE
      )

      if (!inherits(final_model, "try-error")) {
        #sel_names <- selected_predictor_names(final_model, pred_ord)
        #n_selected <- length(sel_names)
        #selected_str <- if (length(sel_names) > 0) paste(sel_names, collapse = ",") else ""

        tcst <- try(safe_predict_probs(final_model, dataF), silent = TRUE)
        
        if (!inherits(tcst, "try-error")) {
          #fcst <- tcst[nrow(tcst), ]
          #below  <- 100 * fcst[1]
          #normal <- 100 * fcst[2]
          #above  <- 100 * fcst[3]

          ntest <- min(length(tdry), nrow(tcst))
          tcst_test <- tcst[seq_len(ntest), , drop = FALSE]

          # save test-period probabilities as row-wise vector:
          # year1(dry,normal,wet), year2(...), ...
          test_probs_vec[seq_len(3 * ntest)] <- as.numeric(t(tcst_test))
	  
	  test_obs_num <- (tdry + 2 * tavg + 3 * twet)[seq_len(ntest)]
          tfcst <- cbind(test_obs_num, tcst_test)

          oktr  <- row_complete(tfcst)
          tfcst <- tfcst[oktr, , drop = FALSE]

          if (nrow(tfcst) > 3) {
            trpss <- zrpss(tfcst)
          }
        }
      }
    } # if !inherit after simlMod

######
### Now use the entire data to redevelop model for final forecast
#####
    DOTHIS <- TRUE
    if (DOTHIS) {
       simlMod <- try(multinom(lower_formula, data = dataHz, trace = FALSE), silent = TRUE)
       if (!inherits(simlMod, "try-error")) {
          final_model <- try(
          stepAIC(simlMod,
                  scope = list(lower = lower_formula, upper = upper_formula),
                  direction = "forward",
                  trace = 0),
          silent = TRUE
          )

         if (!inherits(final_model, "try-error")) {
           sel_names <- selected_predictor_names(final_model, pred_ord)
           n_selected <- length(sel_names)
           selected_str <- if (length(sel_names) > 0) paste(sel_names, collapse = ",") else ""

           tcst <- try(safe_predict_probs(final_model, dataFz), silent = TRUE)

           if (!inherits(tcst, "try-error")) {
             fcst <- tcst[nrow(tcst), ]
             below  <- 100 * fcst[1]
             normal <- 100 * fcst[2]
             above  <- 100 * fcst[3]

           }
         }
       } # if !inherit after simlMod
    } # DOTHIS
######

  } else {  # else of Leave3Out

    simlMod <- try(multinom(lower_formula, data = dataH, trace = FALSE), silent = TRUE)

    if (!inherits(simlMod, "try-error")) {
      final_model <- try(
        stepAIC(simlMod,
                scope = list(lower = lower_formula, upper = upper_formula),
                direction = "forward",
                trace = 0),
        silent = TRUE
      )

      if (!inherits(final_model, "try-error")) {
        sel_names <- selected_predictor_names(final_model, pred_ord)
        n_selected <- length(sel_names)
        selected_str <- if (length(sel_names) > 0) paste(sel_names, collapse = ",") else ""

        hcst <- try(safe_predict_probs(final_model, dataH), silent = TRUE)
        if (!inherits(hcst, "try-error")) {
          hc_in <- cbind(obs_num_ok, hcst)
          okhc <- row_complete(hc_in)
          hc_in <- hc_in[okhc, , drop = FALSE]

          if (nrow(hc_in) > 3) {
            hrpss <- zrpss(hc_in)
            hhss  <- safe_hss(hc_in[, 1], hc_in[, 2:4, drop = FALSE], xpb = xpb)
            tmp <- hc_in[,2:4]
            hsharpness <- mean_sharpness(tmp) 
          }
        }

        tcst <- try(safe_predict_probs(final_model, dataF), silent = TRUE)
        if (!inherits(tcst, "try-error")) {
          fcst <- tcst[nrow(tcst), ]

          below  <- 100 * fcst[1]
          normal <- 100 * fcst[2]
          above  <- 100 * fcst[3]

          ntest <- min(length(tdry), nrow(tcst))
          tcst_test <- tcst[seq_len(ntest), , drop = FALSE]

          # save test-period probabilities as row-wise vector:
          # year1(dry,normal,wet), year2(...), ...
          test_probs_vec[seq_len(3 * ntest)] <- as.numeric(t(tcst_test))

	  test_obs_num <- (tdry + 2 * tavg + 3 * twet)[seq_len(ntest)]
          tfcst <- cbind(test_obs_num, tcst_test)

          oktr  <- row_complete(tfcst)
          tfcst <- tfcst[oktr, , drop = FALSE]

          if (nrow(tfcst) > 3) {
            trpss <- zrpss(tfcst)
          }
        }
      }
    }
  }

  score_vec <- c(below, normal, above, hrpss, trpss, hhss, hsharpness)
  score_vec <- as.numeric(score_vec)
  names(score_vec) <- c("below", "normal", "above", "hrpss", "trpss", "hhss","hsharpness")

  ret <- list(
    scores = score_vec,
    n_selected = n_selected,
    selected_str = selected_str,
    test_probs = test_probs_vec
  )

  return(ret)

} # End cell_worker

###########################################################
#####  run_serial_cell
#######################################
run_serial_cells <- function(oblw, onrm, oabv,
                             tblw, tnrm, tabv,
                             zblw, znrm, zabv,
                             Xh, Xf, zXh, zXf,
                             nlon, nlat,
                             nmiss, xpb, Leave3Out,
                             leave_idx,
                             pred_ord,
                             RUN_TYPE) {

  out_names <- c("below", "normal", "above", "hrpss", "trpss", "hhss","hsharpness")
  ncell <- nlon * nlat

  ntest <- dim(tblw)[3]
  test_probs_mat <- matrix(NA_real_, nrow = ncell, ncol = 3 * ntest)

  resmat <- matrix(NA_real_, nrow = ncell, ncol = length(out_names))
  colnames(resmat) <- out_names

  nsel_vec <- rep(NA_real_, ncell)
  sel_str_vec <- rep(NA_character_, ncell)

  for (gid in seq_len(ncell)) {
    ans <- try(
      cell_worker(
        cell_id = gid,
        nlon = nlon, nlat = nlat,
        oblw = oblw, onrm = onrm, oabv = oabv,
        tblw = tblw, tnrm = tnrm, tabv = tabv,
        zblw = zblw, znrm = znrm, zabv = zabv,
        Xh   = Xh,   Xf   = Xf,   zXh  = zXh, zXf = zXf,
        nmiss = nmiss, xpb = xpb, Leave3Out = Leave3Out,
        leave_idx = leave_idx,
        pred_ord = pred_ord,
        RUN_TYPE = RUN_TYPE
      ),
      silent = TRUE
    )

    if (inherits(ans, "try-error") || is.null(ans) || !is.list(ans) || is.null(ans$scores)) {
      resmat[gid, ] <- rep(NA_real_, length(out_names))
      nsel_vec[gid] <- NA_real_
      sel_str_vec[gid] <- NA_character_
      test_probs_mat[gid, ] <- rep(NA_real_, 3 * ntest)
    } else {
      tmp_scores <- ans$scores
      if (is.null(tmp_scores) || length(tmp_scores) != length(out_names)) {
        resmat[gid, ] <- rep(NA_real_, length(out_names))
      } else {
        resmat[gid, ] <- as.numeric(tmp_scores)
      }

      nsel_vec[gid] <- ans$n_selected
      sel_str_vec[gid] <- ans$selected_str

      if (!is.null(ans$test_probs) && length(ans$test_probs) == 3 * ntest) {
        test_probs_mat[gid, ] <- ans$test_probs
      } else {
        test_probs_mat[gid, ] <- rep(NA_real_, 3 * ntest)
      }
    }
    if (gid %% 1000 == 0) {
      cat("  processed", gid, "of", ncell, "grid cells\n")
    }
  }

    test_probs_arr <- array(NA_real_, dim = c(nlon, nlat, ntest, 3))
    for (gid in seq_len(ncell)) {
      ilon <- ((gid - 1) %% nlon) + 1
      jlat <- ((gid - 1) %/% nlon) + 1
      test_probs_arr[ilon, jlat, , ] <- matrix(test_probs_mat[gid, ], nrow = ntest, ncol = 3, byrow = TRUE)
    }

  list(
    below = matrix(resmat[, 1], nrow = nlon, ncol = nlat),
    normal = matrix(resmat[, 2], nrow = nlon, ncol = nlat),
    above = matrix(resmat[, 3], nrow = nlon, ncol = nlat),
    hrpss = matrix(resmat[, 4], nrow = nlon, ncol = nlat),
    trpss = matrix(resmat[, 5], nrow = nlon, ncol = nlat),
    hhss = matrix(resmat[, 6], nrow = nlon, ncol = nlat),
    hsharpness = matrix(resmat[, 7], nrow = nlon, ncol = nlat),
    n_selected = matrix(nsel_vec, nrow = nlon, ncol = nlat),
    selected_str = matrix(sel_str_vec, nrow = nlon, ncol = nlat),
    test_probs = test_probs_arr
  )

} # run_serial_cells

###########################################################
#####  run_parallel_cell
#######################################
run_parallel_cells <- function(oblw, onrm, oabv,
                               tblw, tnrm, tabv,
                               zblw, znrm, zabv,
                               Xh, Xf, zXh, zXf,
                               nlon, nlat,
                               nmiss, xpb, Leave3Out,
                               leave_idx,
                               pred_ord,
                               RUN_TYPE,
                               ncores = 1,
                               use_mclapply = TRUE) {

  out_names <- c("below", "normal", "above", "hrpss", "trpss", "hhss","hsharpness")
  ncell <- nlon * nlat
  grid_ids <- seq_len(ncell)

  ntest <- dim(tblw)[3]
  test_probs_mat <- matrix(NA_real_, nrow = ncell, ncol = 3 * ntest)

  safe_worker <- function(gid) {
    ans <- try(
      cell_worker(
        cell_id = gid,
        nlon = nlon, nlat = nlat,
        oblw = oblw, onrm = onrm, oabv = oabv,
        tblw = tblw, tnrm = tnrm, tabv = tabv,
        zblw = zblw, znrm = znrm, zabv = zabv,
        Xh   = Xh,   Xf   = Xf,   zXh  = zXh, zXf = zXf,
        nmiss = nmiss, xpb = xpb, Leave3Out = Leave3Out,
        leave_idx = leave_idx,
        pred_ord = pred_ord,
        RUN_TYPE = RUN_TYPE
      ),
      silent = TRUE
    )

    if (inherits(ans, "try-error") || is.null(ans) || !is.list(ans) || is.null(ans$scores)) {
      return(list(
        scores = setNames(rep(NA_real_, length(out_names)), out_names),
        n_selected = NA_real_,
        selected_str = NA_character_,
        test_probs = rep(NA_real_, 3 * ntest)
      ))
    }
    ans
  }

  if (use_mclapply) {
    res <- parallel::mclapply(grid_ids, safe_worker, mc.cores = ncores)
  } else {
    cl <- parallel::makeCluster(ncores)
    on.exit(parallel::stopCluster(cl), add = TRUE)

    parallel::clusterEvalQ(cl, {
      library(nnet)
      library(MASS)
      NULL
    })

    parallel::clusterExport(
      cl,
      varlist = c(
        "oblw","onrm","oabv","tblw","tnrm","tabv","zblw","znrm","zabv",
        "Xh","Xf", "zXh","zXf","nlon","nlat","nmiss","xpb","Leave3Out","leave_idx","pred_ord","RUN_TYPE","ntest",
        "obs_cat_from_binary","obs_num_from_binary","row_complete",
        "safe_predict_probs","safe_hss","selected_predictor_names",
        "zrpss","multi.cont","cell_worker","safe_worker","mean_sharpness"
      ),
      envir = environment()
    )

    res <- parallel::parLapply(cl, grid_ids, safe_worker)
  }

  resmat <- matrix(NA_real_, nrow = ncell, ncol = length(out_names))
  nsel_vec <- rep(NA_real_, ncell)
  sel_str_vec <- rep(NA_character_, ncell)

  for (gid in seq_len(ncell)) {
    tmp_scores <- res[[gid]]$scores
    if (is.null(tmp_scores) || length(tmp_scores) != length(out_names)) {
      resmat[gid, ] <- rep(NA_real_, length(out_names))
    } else {
      resmat[gid, ] <- as.numeric(tmp_scores)
    }
    nsel_vec[gid] <- res[[gid]]$n_selected
    sel_str_vec[gid] <- res[[gid]]$selected_str
    if (!is.null(res[[gid]]$test_probs) && length(res[[gid]]$test_probs) == 3 * ntest) {
      test_probs_mat[gid, ] <- res[[gid]]$test_probs
    } else {
      test_probs_mat[gid, ] <- rep(NA_real_, 3 * ntest)
    }  
  }

  test_probs_arr <- array(NA_real_, dim = c(nlon, nlat, ntest, 3))

  for (gid in seq_len(ncell)) {
    ilon <- ((gid - 1) %% nlon) + 1
    jlat <- ((gid - 1) %/% nlon) + 1
    test_probs_arr[ilon, jlat, , ] <-
    matrix(test_probs_mat[gid, ], nrow = ntest, ncol = 3, byrow = TRUE)
  }
  list(
    below = matrix(resmat[, 1], nrow = nlon, ncol = nlat),
    normal = matrix(resmat[, 2], nrow = nlon, ncol = nlat),
    above = matrix(resmat[, 3], nrow = nlon, ncol = nlat),
    hrpss = matrix(resmat[, 4], nrow = nlon, ncol = nlat),
    trpss = matrix(resmat[, 5], nrow = nlon, ncol = nlat),
    hhss = matrix(resmat[, 6], nrow = nlon, ncol = nlat),
    hsharpness = matrix(resmat[, 7], nrow = nlon, ncol = nlat),
    n_selected = matrix(nsel_vec, nrow = nlon, ncol = nlat),
    selected_str = matrix(sel_str_vec, nrow = nlon, ncol = nlat),
    test_probs = test_probs_arr
  )

} # run_parallel_cells

##########################################################
#### Part4: Red Observations, Assign Variables, Compute, and Write
#########################################################
###############################################################################
# READ OBSERVATIONS
###############################################################################
nco <- open.nc(paste0(hinddiro, fileprecnc))
timeo <- var.get.nc(nco, "time")
htimeatt <- att.get.nc(nco, "time", "units")
zodate <- utcal.nc(htimeatt, timeo, type = "n")

preco <- var.get.nc(nco, "prec")
lon   <- var.get.nc(nco, "lon")
lat   <- var.get.nc(nco, "lat")

lon_chr <- as.character(lon)
lat_chr <- as.character(lat)

oDateNames <- paste0(zodate[, 1], SEASON)
names(oDateNames) <- oDateNames
dimnames(preco) <- list(lon_chr, lat_chr, oDateNames)
preco[preco < 0] <- NA

###############################################################################
# DETERMINE TRAIN/TEST YEARS
###############################################################################
nch0 <- open.nc(paste0(hinddir, hindcastNames[1]))
tunitfile <- att.get.nc(nch0, "time", "units")
timeh0 <- var.get.nc(nch0, "time")
zhdate <- utcal.nc(tunitfile, timeh0, type = "n")
close.nc(nch0)

obsyrs  <- zodate[, 1]
hindyrs <- zhdate[, 1]

lead_to_select <- seq(iLead, iLead + SznLen - 1)

#obs_month_to_select <- seq(iMonth + iLead, iMonth + iLead + SznLen - 1)
#tmp <- month.2abb[obs_month_to_select]

#if (SznLen == 4) {
#  SEASON1 <- paste0(substr(tmp[1],1,1), substr(tmp[2],1,1), substr(tmp[3],1,1), substr(tmp[4],1,1))
#} else {
#  SEASON1 <- paste0(substr(tmp[1],1,1), substr(tmp[2],1,1), substr(tmp[3],1,1))
#}

#if (SEASON != SEASON1) stop(paste("WRONG", SEASON))

#obs_month_to_select <- obs_month_to_select %% 12
#zeroind <- which(obs_month_to_select == 0)
#if (length(zeroind) > 0) obs_month_to_select[zeroind] <- 12

# 1. Generate the sequence of months (which might go past 12)
raw_months <- seq(iMonth + iLead, iMonth + iLead + SznLen - 1)

# 2. Use modular arithmetic to safely wrap months back to 1-12 (e.g., 13 -> 1)
obs_month_to_select <- ((raw_months - 1) %% 12) + 1
tmp <- month.abb[obs_month_to_select]

# 3. Create the season name (initials for seasons, full abbreviation for single months)
SEASON1 <- if (SznLen > 1) paste0(substr(tmp, 1, 1), collapse = "") else tmp

# 4. Validation check
if (SEASON != SEASON1) {
  stop(sprintf("Season mismatch! Expected: %s, Generated: %s", SEASON, SEASON1))
}

lastMonth  <- obs_month_to_select[length(obs_month_to_select)]
firstMonth <- obs_month_to_select[1]

lastObsYr <- LastYr

lastMonthIndex <- which(obsyrs == lastObsYr)
obsyrs <- obsyrs[1:lastMonthIndex]

is_next_year_season <- (iMonth + iLead) > 12
year_shift <- if (is_next_year_season) 1 else 0

hind_target_years  <- hindyrs + year_shift
cyrs               <- intersect(obsyrs, hind_target_years)
original_hind_years <- cyrs - year_shift
yrFCST <- hind_target_years[length(hind_target_years)]

totalTrainYrs <- length(cyrs)
trainLength <- totalTrainYrs - testLength
trainYrsObs <- cyrs[1:trainLength]
testYrsObs  <- cyrs[(trainLength+1):length(cyrs)]

trainYrsMod <- original_hind_years[1:trainLength]
testYrsMod  <- original_hind_years[(trainLength+1):length(cyrs)]

trainObs <- paste0(trainYrsObs, SEASON)
testObs  <- paste0(testYrsObs, SEASON)

trainMod <- paste0(trainYrsMod, SEASON)
testMod  <- paste0(testYrsMod, SEASON)

fcstyrsMod <- union(testYrsMod, LastYrHind)
fcstyrsObs <- union(testYrsObs, yrFCST)

fcstMod <- paste0(fcstyrsMod, SEASON)
fcstObs <- paste0(fcstyrsObs, SEASON)

allyrsMod <- c(trainMod, fcstMod)
allyrsObs <- c(trainObs, fcstObs)

###############################################################################
# PREPARE OBSERVATION CATEGORIES
###############################################################################
yobs <- preco[, , trainObs, drop = FALSE]
tOBS <- preco[, , testObs,  drop = FALSE]
zOBS <- preco[, , c(trainObs,testObs),  drop = FALSE]

dimnames(yobs) <- list(lon_chr, lat_chr, trainObs)
dimnames(zOBS) <- list(lon_chr, lat_chr, c(trainObs,testObs))

yobs[yobs < rthr] <- NA
yobs <- round(yobs, 1)

zOBS[zOBS < rthr ] <- NA
zOBS <- round(zOBS,1)

hblw <- apply(yobs, c(1, 2), quantile, probs = 1/3, na.rm = TRUE)
habv <- apply(yobs, c(1, 2), quantile, probs = 2/3, na.rm = TRUE)

zhblw <- apply(zOBS, c(1, 2), quantile, probs = 1/3, na.rm = TRUE)
zhabv <- apply(zOBS, c(1, 2), quantile, probs = 2/3, na.rm = TRUE)

ntot <- dim(yobs)[3]
hblw <- round(array(rep(hblw, ntot), c(dim(hblw), ntot)), 1)
habv <- round(array(rep(habv, ntot), c(dim(habv), ntot)), 1)

zntot <- dim(zOBS)[3]
zhblw <- round(array(rep(zhblw, zntot), c(dim(zhblw), zntot)), 1)
zhabv <- round(array(rep(zhabv, zntot), c(dim(zhabv), zntot)), 1)

oblw <- 1 * (yobs < hblw)
oabv <- 1 * (yobs > habv)
onrm <- 1 * (yobs >= hblw & yobs <= habv)

zblw <- 1 * (zOBS < zhblw)
zabv <- 1 * (zOBS > zhabv)
znrm <- 1 * (zOBS >= zhblw & zOBS <= zhabv)

tOBS[tOBS < rthr] <- NA
tOBS <- round(tOBS, 1)

tmp1 <- hblw[, , 1:dim(tOBS)[3], drop = FALSE]
tmp2 <- habv[, , 1:dim(tOBS)[3], drop = FALSE]

tblw <- 1 * (tOBS < tmp1)
tabv <- 1 * (tOBS > tmp2)
tnrm <- 1 * (tOBS >= tmp1 & tOBS <= tmp2)

if (stdizeObs) {
  y <- yobs
  ysum <- apply(y, c(1,2), sum, na.rm = TRUE)
  ysum[ysum < rthr] <- NA

  ycnt <- 1 * (!is.na(y))
  ycnt <- apply(ycnt, c(1,2), sum)
  ycnt[ycnt <= iselLen] <- NA

  oMean <- ysum / ycnt
  oMean <- array(rep(oMean, ntot), c(dim(oMean), ntot))
  oAnom <- yobs - oMean

  ycnt2 <- ycnt - 1
  anomsum <- apply(oAnom * oAnom, c(1,2), sum, na.rm = TRUE)
  oSdev <- sqrt(anomsum / ycnt2)
  oSdev[oSdev < 0.1] <- NA
  oSdev <- array(rep(oSdev, ntot), c(dim(oSdev), ntot))

  oZscore <- oAnom / oSdev

  hblw <- apply(oZscore, c(1,2), quantile, probs = 1/3, na.rm = TRUE)
  habv <- apply(oZscore, c(1,2), quantile, probs = 2/3, na.rm = TRUE)
  hblw <- round(array(rep(hblw, ntot), c(dim(hblw), ntot)), 1)
  habv <- round(array(rep(habv, ntot), c(dim(habv), ntot)), 1)

  yobs <- oZscore

  oblw <- 1 * (yobs < hblw)
  oabv <- 1 * (yobs > habv)
  onrm <- 1 * (yobs >= hblw & yobs <= habv)

  ### for zOBS
  y <- zOBS
  ysum <- apply(y, c(1,2), sum, na.rm = TRUE)
  ysum[ysum < rthr] <- NA

  ycnt <- 1 * (!is.na(y))
  ycnt <- apply(ycnt, c(1,2), sum)
  ycnt[ycnt <= iselLen] <- NA

  zoMean <- ysum / ycnt
  zoMean <- array(rep(zoMean, zntot), c(dim(zoMean), zntot))
  zoAnom <- zOBS - zoMean

  ycnt2 <- ycnt - 1
  anomsum <- apply(zoAnom * zoAnom, c(1,2), sum, na.rm = TRUE)
  zoSdev <- sqrt(anomsum / ycnt2)
  zoSdev[zoSdev < 0.1] <- NA
  zoSdev <- array(rep(zoSdev, zntot), c(dim(zoSdev), zntot))

  zoZscore <- zoAnom / zoSdev

  zhblw <- apply(zoZscore, c(1,2), quantile, probs = 1/3, na.rm = TRUE)
  zhabv <- apply(zoZscore, c(1,2), quantile, probs = 2/3, na.rm = TRUE)
  zhblw <- round(array(rep(zhblw, zntot), c(dim(zhblw), zntot)), 1)
  zhabv <- round(array(rep(zhabv, zntot), c(dim(zhabv), zntot)), 1)

  zOBS <- zoZscore

  zblw <- 1 * (zOBS < zhblw)
  zabv <- 1 * (zOBS > zhabv)
  znrm <- 1 * (zOBS >= zhblw & zOBS <= zhabv)


  ## for tOBS
  
  tmp1 <- oMean[, , 1:dim(tOBS)[3], drop = FALSE]
  tmp2 <- oSdev[, , 1:dim(tOBS)[3], drop = FALSE]
  tOBS <- (tOBS - tmp1) / tmp2

  tmp1 <- hblw[, , 1:dim(tOBS)[3], drop = FALSE]
  tmp2 <- habv[, , 1:dim(tOBS)[3], drop = FALSE]

  tblw <- 1 * (tOBS < tmp1)
  tabv <- 1 * (tOBS > tmp2)
  tnrm <- 1 * (tOBS >= tmp1 & tOBS <= tmp2)
}

###############################################################################
# DIAGNOSTIC GRID CELLS (SAFE – AFTER oblw IS DEFINED)
###############################################################################
if (RUN_DIAGNOSTIC) {
  diag_cells <- list(
    list(
      ilon = max(1, floor(dim(oblw)[1] / 2)),
      jlat = max(1, floor(dim(oblw)[2] / 2))
    )
  )
}
###############################################################################
# READ ALL MODEL DATA
###############################################################################
wnpAll  <- matrix(NA, nrow = length(allyrsMod), ncol = length(forecastNames), dimnames = list(allyrsMod, forecastNames))
wepAll  <- wnpAll
ni4All  <- wnpAll
n34All  <- wnpAll
wvgAll  <- wnpAll
mgiAll  <- wnpAll
atl3All <- wnpAll
iodAll  <- wnpAll
siodAll <- wnpAll
rn34All <- wnpAll

pc1All  <- wnpAll
pc2All  <- wnpAll
pc3All  <- wnpAll
pc4All  <- wnpAll
pc5All  <- wnpAll
pc6All  <- wnpAll
rc1All  <- wnpAll
rc2All  <- wnpAll
rc3All  <- wnpAll
rc4All  <- wnpAll
rc5All  <- wnpAll
rc6All  <- wnpAll

for (imodel in gcmStart:gcmLast) {
  ncf <- open.nc(paste0(fcstdir, hindcastNames[imodel]))

  tunitfile <- att.get.nc(ncf, "time", "units")
  timeh <- var.get.nc(ncf, "time")
  zhdate <- utcal.nc(tunitfile, timeh, type = "n")
  yrs <- paste0(zhdate[,1], SEASON)

  yrsind <- 1:length(yrs)
  names(yrsind) <- yrs
  yrs <- intersect(yrs, allyrsMod)
  ind.used <- yrsind[yrs]

  wnp  <- var.get.nc(ncf, "wnp")
  wep  <- var.get.nc(ncf, "wep")
  ni4  <- var.get.nc(ncf, "ni4")
  n34  <- var.get.nc(ncf, "n34")
  wvg  <- var.get.nc(ncf, "wvg")
  mgi  <- var.get.nc(ncf, "mgi")
  atl3 <- var.get.nc(ncf, "atn")
  iod  <- var.get.nc(ncf, "iod")
  siod <- var.get.nc(ncf, "siod")
  rn34 <- var.get.nc(ncf, "rn34")
  close.nc(ncf)

  allpc <- read.table(paste0(hinddir, eoftsNames[imodel]))
  allrc <- read.table(paste0(hinddir, eofPrecNames[imodel]))

  wnpAll[yrs, imodel]  <- wnp[ind.used]
  wepAll[yrs, imodel]  <- wep[ind.used]
  ni4All[yrs, imodel]  <- ni4[ind.used]
  n34All[yrs, imodel]  <- n34[ind.used]
  wvgAll[yrs, imodel]  <- wvg[ind.used]
  mgiAll[yrs, imodel]  <- mgi[ind.used]
  atl3All[yrs, imodel] <- atl3[ind.used]
  iodAll[yrs, imodel]  <- iod[ind.used]
  siodAll[yrs, imodel] <- siod[ind.used]
  rn34All[yrs, imodel] <- rn34[ind.used]

  pc1All[yrs, imodel]  <- allpc[ind.used,1]
  pc2All[yrs, imodel]  <- allpc[ind.used,2]
  pc3All[yrs, imodel]  <- allpc[ind.used,3]
  pc4All[yrs, imodel]  <- allpc[ind.used,4]
  pc5All[yrs, imodel]  <- allpc[ind.used,5]
  pc6All[yrs, imodel]  <- allpc[ind.used,6]
  rc1All[yrs, imodel]  <- allrc[ind.used,1]
  rc2All[yrs, imodel]  <- allrc[ind.used,2]
  rc3All[yrs, imodel]  <- allrc[ind.used,3]
  rc4All[yrs, imodel]  <- allrc[ind.used,4]
  rc5All[yrs, imodel]  <- allrc[ind.used,5]
  rc6All[yrs, imodel]  <- allrc[ind.used,6]

}

###############################################################################
# SELECT TRAIN/FCST YEARS
###############################################################################
wnph  <- wnpAll[trainMod, gcmStart:gcmLast, drop = FALSE]
weph  <- wepAll[trainMod, gcmStart:gcmLast, drop = FALSE]
ni4h  <- ni4All[trainMod, gcmStart:gcmLast, drop = FALSE]
n34h  <- n34All[trainMod, gcmStart:gcmLast, drop = FALSE]
wvgh  <- wvgAll[trainMod, gcmStart:gcmLast, drop = FALSE]
mgih  <- mgiAll[trainMod, gcmStart:gcmLast, drop = FALSE]
atl3h <- atl3All[trainMod, gcmStart:gcmLast, drop = FALSE]
iodh  <- iodAll[trainMod, gcmStart:gcmLast, drop = FALSE]
siodh <- siodAll[trainMod, gcmStart:gcmLast, drop = FALSE]
rn34h <- rn34All[trainMod, gcmStart:gcmLast, drop = FALSE]

pc1h  <- pc1All[trainMod, gcmStart:gcmLast, drop = FALSE]
pc2h  <- pc2All[trainMod, gcmStart:gcmLast, drop = FALSE]
pc3h  <- pc3All[trainMod, gcmStart:gcmLast, drop = FALSE]
pc4h  <- pc4All[trainMod, gcmStart:gcmLast, drop = FALSE]
pc5h  <- pc5All[trainMod, gcmStart:gcmLast, drop = FALSE]
pc6h  <- pc6All[trainMod, gcmStart:gcmLast, drop = FALSE]

rc1h  <- rc1All[trainMod, gcmStart:gcmLast, drop = FALSE]
rc2h  <- rc2All[trainMod, gcmStart:gcmLast, drop = FALSE]
rc3h  <- rc3All[trainMod, gcmStart:gcmLast, drop = FALSE]
rc4h  <- rc4All[trainMod, gcmStart:gcmLast, drop = FALSE]
rc5h  <- rc5All[trainMod, gcmStart:gcmLast, drop = FALSE]
rc6h  <- rc6All[trainMod, gcmStart:gcmLast, drop = FALSE]

wnpf  <- wnpAll[fcstMod, gcmStart:gcmLast, drop = FALSE]
wepf  <- wepAll[fcstMod, gcmStart:gcmLast, drop = FALSE]
ni4f  <- ni4All[fcstMod, gcmStart:gcmLast, drop = FALSE]
n34f  <- n34All[fcstMod, gcmStart:gcmLast, drop = FALSE]
wvgf  <- wvgAll[fcstMod, gcmStart:gcmLast, drop = FALSE]
mgif  <- mgiAll[fcstMod, gcmStart:gcmLast, drop = FALSE]
atl3f <- atl3All[fcstMod, gcmStart:gcmLast, drop = FALSE]
iodf  <- mgiAll[fcstMod, gcmStart:gcmLast, drop = FALSE]
siodf <- siodAll[fcstMod, gcmStart:gcmLast, drop = FALSE]
rn34f <- rn34All[fcstMod, gcmStart:gcmLast, drop = FALSE]

pc1f  <- pc1All[fcstMod, gcmStart:gcmLast, drop = FALSE]
pc2f  <- pc2All[fcstMod, gcmStart:gcmLast, drop = FALSE]
pc3f  <- pc3All[fcstMod, gcmStart:gcmLast, drop = FALSE]
pc4f  <- pc4All[fcstMod, gcmStart:gcmLast, drop = FALSE]
pc5f  <- pc5All[fcstMod, gcmStart:gcmLast, drop = FALSE]
pc6f  <- pc6All[fcstMod, gcmStart:gcmLast, drop = FALSE]

rc1f  <- rc1All[fcstMod, gcmStart:gcmLast, drop = FALSE]
rc2f  <- rc2All[fcstMod, gcmStart:gcmLast, drop = FALSE]
rc3f  <- rc3All[fcstMod, gcmStart:gcmLast, drop = FALSE]
rc4f  <- rc4All[fcstMod, gcmStart:gcmLast, drop = FALSE]
rc5f  <- rc5All[fcstMod, gcmStart:gcmLast, drop = FALSE]
rc6f  <- rc6All[fcstMod, gcmStart:gcmLast, drop = FALSE]

###############################################################################
# OPTIONAL DETREND
###############################################################################
if (DETREND) {
  ones <- seq_along(trainMod)
  tyr  <- fcstyrsMod

  detrend_one <- function(hmat, fmat, trainYrsMod, tyr) {
    h2 <- hmat
    f2 <- fmat
    for (im in seq_len(ncol(hmat))) {
      fit <- lm(h2[, im] ~ ones)
      fv  <- fit$fitted.values

      ok <- which(is.finite(h2[, im]) & is.finite(fv))
      h2[ok, im] <- h2[ok, im] - fv[ok]
      f2[, im]   <- f2[, im] - (fit$coef[1] + fit$coef[2] * (tyr - trainYrsMod[1] + 1))
    }
    list(h = h2, f = f2)
  }

  z <- detrend_one(wnph,  wnpf,  trainYrsMod, tyr); wnph  <- z$h; wnpf  <- z$f
  z <- detrend_one(weph,  wepf,  trainYrsMod, tyr); weph  <- z$h; wepf  <- z$f
  z <- detrend_one(ni4h,  ni4f,  trainYrsMod, tyr); ni4h  <- z$h; ni4f  <- z$f
  z <- detrend_one(n34h,  n34f,  trainYrsMod, tyr); n34h  <- z$h; n34f  <- z$f
  z <- detrend_one(wvgh,  wvgf,  trainYrsMod, tyr); wvgh  <- z$h; wvgf  <- z$f
  z <- detrend_one(mgih,  mgif,  trainYrsMod, tyr); mgih  <- z$h; mgif  <- z$f
  z <- detrend_one(atl3h, atl3f, trainYrsMod, tyr); atl3h <- z$h; atl3f <- z$f
  z <- detrend_one(iodh,  iodf,  trainYrsMod, tyr); iodh  <- z$h; iodf  <- z$f
  z <- detrend_one(siodh, siodf, trainYrsMod, tyr); siodh <- z$h; siodf <- z$f
  z <- detrend_one(rn34h, rn34f, trainYrsMod, tyr); rn34h <- z$h; rn34f <- z$f

  z <- detrend_one(pc1h,  pc1f,  trainYrsMod, tyr); pc1h  <- z$h; pc1f  <- z$f
  z <- detrend_one(pc2h,  pc2f,  trainYrsMod, tyr); pc2h  <- z$h; pc2f  <- z$f
  z <- detrend_one(pc3h,  pc3f,  trainYrsMod, tyr); pc3h  <- z$h; pc3f  <- z$f
  z <- detrend_one(pc4h,  pc4f,  trainYrsMod, tyr); pc4h  <- z$h; pc4f  <- z$f
  z <- detrend_one(pc5h,  pc5f,  trainYrsMod, tyr); pc5h  <- z$h; pc5f  <- z$f
  z <- detrend_one(pc6h,  pc6f,  trainYrsMod, tyr); pc6h  <- z$h; pc6f  <- z$f

  z <- detrend_one(rc1h,  rc1f,  trainYrsMod, tyr); rc1h  <- z$h; rc1f  <- z$f
  z <- detrend_one(rc2h,  rc2f,  trainYrsMod, tyr); rc2h  <- z$h; rc2f  <- z$f
  z <- detrend_one(rc3h,  rc3f,  trainYrsMod, tyr); rc3h  <- z$h; rc3f  <- z$f
  z <- detrend_one(rc4h,  rc4f,  trainYrsMod, tyr); rc4h  <- z$h; rc4f  <- z$f
  z <- detrend_one(rc5h,  rc5f,  trainYrsMod, tyr); rc5h  <- z$h; rc5f  <- z$f
  z <- detrend_one(rc6h,  rc6f,  trainYrsMod, tyr); rc6h  <- z$h; rc6f  <- z$f

}

###############################################################################
# ASSIGN Z-NAME OBJECTS FOR COMPATIBILITY
###############################################################################
wnphZ  <- wnph;  wephZ  <- weph;  ni4hZ <- ni4h;  Nino34hZ <- n34h
WVGhZ  <- wvgh;  MGIhZ  <- mgih;  ATL3hZ <- atl3h
IODhZ  <- iodh;  SIODhZ <- siodh; rNino34hZ  <- rn34h

PCS1hZ  <- pc1h;  PCS2hZ  <- pc2h;  PCS3hZ <- pc3h
PCS4hZ  <- pc4h;  PCS5hZ  <- pc5h;  PCS6hZ <- pc6h

PCR1hZ  <- rc1h;  PCR2hZ  <- rc2h;  PCR3hZ <- rc3h
PCR4hZ  <- rc4h;  PCR5hZ  <- rc5h;  PCR6hZ <- rc6h

wnpfZ  <- wnpf;  wepfZ  <- wepf;  ni4fZ <- ni4f;  Nino34fZ <- n34f
WVGfZ  <- wvgf;  MGIfZ  <- mgif;  ATL3fZ <- atl3f
IODfZ  <- iodf;  SIODfZ <- siodf; rNino34fZ <- rn34f

PCS1fZ  <- pc1f;  PCS2fZ  <- pc2f;  PCS3fZ <- pc3f
PCS4fZ  <- pc4f;  PCS5fZ  <- pc5f;  PCS6fZ <- pc6f

PCR1fZ  <- rc1f;  PCR2fZ  <- rc2f;  PCR3fZ <- rc3f
PCR4fZ  <- rc4f;  PCR5fZ  <- rc5f;  PCR6fZ <- rc6f

hind_list <- mget(AllhNames)
names(hind_list) <- AllPRDCTRS

fcst_list <- mget(AllfNames)
names(fcst_list) <- AllPRDCTRS

###############################################################################
# PRECOMPUTE LEAVE-3-OUT INDICES
###############################################################################
leave_idx_global <- build_leave3_indices(trainLength)


###############################################################################
# RUN CONFIGURATION
###############################################################################
RUN_NAMES <- c(PRDCTRS[seq_len(nprdFirst)], "ALLP_INT")
nrun  <- length(RUN_NAMES)

nlon <- length(lon)
nlat <- length(lat)

ntest <- dim(tblw)[3]

#if (!dir.exists("temp_results")) dir.create("temp_results")

#mylist <-  list(lon = lon, lat = lat, testYrsObs = testYrsObs)
#mypath <- "temp_results/lon_lat_ntest.rds"
#saveRDS(mylist, mypath)

#testcat <- list(tblw = tblw, tnrm = tnrm, tabv = tabv)
#testpath <- "temp_results/testcat.rds"
#saveRDS(testcat, testpath)


hhss_all <- array(NA_real_, dim = c(nlon, nlat, nrun, nmodels))
test_probs_all <- array(NA_real_, dim = c(nlon, nlat, ntest, 3, nrun, nmodels))
hsharpness_all <- array(NA_real_, dim = c(nlon, nlat, nrun, nmodels))

###############################################################################
# MAIN LOOP
###############################################################################
hindcastNames <- hindcastNames[gcmStart:gcmLast]

#task_id <- as.numeric(Sys.getenv("SLURM_ARRAY_TASK_ID", unset = "1"))
#ncores  <- as.numeric(Sys.getenv("SLURM_CPUS_PER_TASK", unset = "40"))

for (irun in seq_len(nrun)) 
{
# Create a table of all 133 combinations (19 runs * 7 models)
# This mimics your nested loop structure exactly
#job_table <- expand.grid(imodel = 1:nmodels, irun = 1:nrun)

# Pick the specific parameters for THIS Slurm task
#current_params <- job_table[task_id, ]
#irun <- current_params$irun
#imodel <- current_params$imodel

  if (irun <= nprdFirst) {
    RUN_TYPE <- "forced_p1"
    run_name <- PRDCTRS[irun]

    notiprd  <- setdiff(seq_len(nprdct), irun)
    pred_ord <- PRDCTRS[c(irun, notiprd)]

  } else {
    RUN_TYPE <- "intercept_all"
    run_name <- "ALLP_INT"

    pred_ord <- PRDCTRS
  }

  for (imodel in 1:nmodels) 
  {

    zproctime1 <- proc.time()
    cat("Starting model", hindcastNames[imodel], " run ", run_name,
        " type ", RUN_TYPE, "\n")

    mats <- make_predictor_matrix(hind_list, fcst_list, pred_ord, imodel)
    Xh <- mats$Xh
    Xf <- mats$Xf
    zXa <- rbind(Xh,Xf)
    zXh  <- zXa[1:zntot,]
    zXf  <- zXa[nrow(zXa), , drop = FALSE]

    leave_idx <- build_leave3_indices(nrow(Xh))

    if (RUN_DIAGNOSTIC) {
     diag_con <- file(diag_file, open = if ((irun == 1) && (imodel == 1)) "wt" else "at")
     on.exit({
       try({ 
        if(inherits(diag_con, "connection") && isOpen(diag_con)) close(diag_con)
       }, silent = TRUE)
     }, add = TRUE)

#     print(getwd())

#     if (RUN_DIAGNOSTIC) {
  	diag_cells <- list(
    	   list(
      	     ilon = max(1, floor(dim(oblw)[1] / 2)),
      	     jlat = max(1, floor(dim(oblw)[2] / 2))
    	   )
  	)
	cat("Diagnostic grid (ilon, jlat):",
        diag_cells[[1]]$ilon, diag_cells[[1]]$jlat, "\n")

     cat("\n\n############################################################\n", file = diag_con)
     cat("MODEL FILE : ", hindcastNames[imodel], "\n", file = diag_con)
     cat("MODEL NAME : ", forecastNames[imodel], "\n", file = diag_con)
     cat("RUN NAME   : ", run_name, "\n", file = diag_con)
     cat("RUN TYPE   : ", RUN_TYPE, "\n", file = diag_con)
     cat("pred_ord   : ", paste(pred_ord, collapse = ","), "\n", file = diag_con)
     cat("############################################################\n", file = diag_con)
     flush(diag_con)

     for (cc in diag_cells) {
      try(
        diagnose_one_cell(
         ilon = cc$ilon,
         jlat = cc$jlat,
         Xh = Xh,
         Xf = Xf,
         pred_ord = pred_ord,
         oblw = oblw, onrm = onrm, oabv = oabv,
         tblw = tblw, tnrm = tnrm, tabv = tabv,
         leave_idx = leave_idx,
         trainYrsObs = trainYrsObs,
         trainYrsMod = trainYrsMod,
         testYrsObs = testYrsObs,
         testYrsMod = testYrsMod,
         fcstyrsObs = fcstyrsObs,
         fcstyrsMod = fcstyrsMod,
         xpb = xpb,
         nmiss = nmiss,
         max_print = diag_max_folds_to_print,
         con = diag_con
        ),
        silent = FALSE
      )
      flush(diag_con)
     }
     #close(diag_con)
     #}
    } # if run Diagnostic

# This bypasses some SSH version checks that might be triggering the SSL error
#    options(parallel.ssh.args = "-o StrictHostKeyChecking=no")

    if (RUN_SERIAL) {
      ans <- run_serial_cells(
        oblw = oblw, onrm = onrm, oabv = oabv,
        tblw = tblw, tnrm = tnrm, tabv = tabv,
        zblw = zblw, znrm = znrm, zabv = zabv,
        Xh = Xh, Xf = Xf, zXh = zXh, zXf = zXf,
        nlon = nlon, nlat = nlat,
        nmiss = nmiss, xpb = xpb, Leave3Out = Leave3Out,
        leave_idx = leave_idx,
        pred_ord = pred_ord,
        RUN_TYPE = RUN_TYPE
      )
    } else {
      ans <- run_parallel_cells(
        oblw = oblw, onrm = onrm, oabv = oabv,
        tblw = tblw, tnrm = tnrm, tabv = tabv,
        zblw = zblw, znrm = znrm, zabv = zabv,
        Xh = Xh, Xf = Xf, zXh = zXh, zXf = zXf,
        nlon = nlon, nlat = nlat,
        nmiss = nmiss, xpb = xpb, Leave3Out = Leave3Out,
        leave_idx = leave_idx,
        pred_ord = pred_ord,
        RUN_TYPE = RUN_TYPE,
        ncores = ncores ,
        use_mclapply = TRUE 
      )
    }

    domain_diag_summary(ans, label = paste(forecastNames[imodel], run_name))

    below        <- ans$below
    normal       <- ans$normal
    above        <- ans$above
    hrpss        <- ans$hrpss
    trpss        <- ans$trpss
    hhss         <- ans$hhss
    hsharpness   <- ans$hsharpness

    hsharpness_all[, , irun, imodel] <- ans$hsharpness
    hhss_all[, , irun, imodel] <- ans$hhss
    test_probs_all[, , , , irun, imodel] <- ans$test_probs

#    save_list <- list(hhss = ans$hhss, test_probs = ans$test_probs, hsharpness = ans$hsharpness)
#    save_path <- sprintf("temp_results/data_irun%02d_mod%02d.rds", irun, imodel)
#    saveRDS(save_list, save_path)

    n_selected   <- ans$n_selected
    selected_str <- ans$selected_str

    dirdat <- paste0(dirdat0, forecastNames[imodel], VARNMTHD)
    fileName <- paste0(run_name, FTYPE)

    source("templateWrite2ncMultinomRgr.R")

    sel_df <- expand.grid(lon = lon, lat = lat)
    sel_df$n_selected   <- as.vector(n_selected)
    sel_df$selected_str <- as.vector(selected_str)

    csvfile <- paste0(dirdat, run_name, FTYPEF, "_selected_predictors.csv")
    write.csv(sel_df, csvfile, row.names = FALSE)

    print(table(n_selected, useNA = "ifany"))

    all_sel <- selected_str[!is.na(selected_str) & nzchar(selected_str)]

    if (length(all_sel) > 0) {
      parts <- unlist(strsplit(all_sel, ",", fixed = TRUE))
      freq_tab <- sort(table(parts), decreasing = TRUE)
      freq_df <- data.frame(
        predictor = names(freq_tab),
        count = as.integer(freq_tab)
      )
      freqfile <- paste0(dirdat, run_name, FTYPEF, "_predictor_frequency.csv")
      write.csv(freq_df, freqfile, row.names = FALSE)
    }


    cat("Done NetCDF for Multinomial Regression\n")
    cat("Completed fcst for imodel =", imodel, " run =", run_name, "\n")
    print(date())

    zproctime2 <- proc.time()
    print(zproctime2 - zproctime1)

  } # imodel
} # irun

######
## Done npredictors and nmodels
## Note that non of the models are in memory as they were run across arrays
## Hence we now have to read saved files
########

DOTHIS <- TRUE
if (DOTHIS) {
  which_max_na <- function(x) {
    if (all(!is.finite(x))) return(NA_integer_)
    which.max(ifelse(is.finite(x), x, -Inf))
  }

best_run_by_model <- array(NA_integer_, dim = c(nlon, nlat, nmodels))
best_hhss_by_model <- array(NA_real_, dim = c(nlon, nlat, nmodels))
best_probs_by_model <- array(NA_real_, dim = c(nlon, nlat, ntest, 3, nmodels))

for (imodel in 1:nmodels) {
  for (ilon in 1:nlon) {
    for (jlat in 1:nlat) {
      hvec <- hhss_all[ilon, jlat, , imodel]
      ir_best <- which_max_na(hvec)

      if (is.finite(ir_best)) {
        best_run_by_model[ilon, jlat, imodel] <- ir_best
        best_hhss_by_model[ilon, jlat, imodel] <- hhss_all[ilon, jlat, ir_best, imodel]
        best_probs_by_model[ilon, jlat, , , imodel] <-
        test_probs_all[ilon, jlat, , , ir_best, imodel]
      }
    }
  }
}

#compute average probabilities for each of the ntest of all models
ens_probs <- apply(best_probs_by_model, c(1, 2, 3, 4), mean, na.rm = TRUE)

#Renormalize
for (ilon in 1:nlon) {
  for (jlat in 1:nlat) {
    for (it in 1:ntest) {
      ps <- ens_probs[ilon, jlat, it, ]
      if (all(is.finite(ps))) {
        s <- sum(ps)
        if (s > 0) ens_probs[ilon, jlat, it, ] <- ps / s
      }
    }
  }
}

test_obs_num_all <- tblw + 2 * tnrm + 3 * tabv

ens_trpss <- matrix(NA_real_, nrow = nlon, ncol = nlat)

for (ilon in 1:nlon) {
  for (jlat in 1:nlat) {
    obs <- test_obs_num_all[ilon, jlat, ]
    pr  <- ens_probs[ilon, jlat, , , drop = FALSE]
    pr  <- matrix(pr, nrow = ntest, ncol = 3)

    tfcst <- cbind(obs, pr)
    ok <- row_complete(tfcst)
    tfcst <- tfcst[ok, , drop = FALSE]

    if (nrow(tfcst) > 3) {
      ens_trpss[ilon, jlat] <- zrpss(tfcst)
    }
  }
}

#close.nc(nco)
test_years <- testYrsObs
source("templateWriteBestModels.R")
source("templateFrequency.R")
source("templateRealTimeHSS.R")

print(date())
zproctime2 <- proc.time()
print(zproctime2 - zproctime0)

}


