
######
## Apr 25-26, 20265-26, 20265-26, 20265-26, 20265-26, 2026
## Functions needed below -- This function used to be with in the main R code
## This could not be done with parallel run, so I needed to make it separate
## A number of functions are sourced here. Write functions are at the end.
########
source("params.R")
source("templateRPS_Wilks.R")
source("templateVerification.R")

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)
}

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

row_complete <- function(x) {
  rowSums(is.finite(x)) == ncol(x)
}

######
## 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
########

AllPRDCTRS <- slctPrdctr

names(AllPRDCTRS) <- AllPRDCTRS

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

TmodelFileNames <- c("canesm5","gem5nemo","nasa","ccsm4","cfsv2","cesm1","ensmean")
AllmodelNamesPlot  <- c("ESM5", "GEMO","NASA","CCSM","CFS2","CESM","ENS")
names(AllmodelNamesPlot) <- TmodelFileNames

modelNamesPlt <- AllmodelNamesPlot[modelFileNames]

nprdctrs <- length(slctPrdctr)
nmodels <- length(modelFileNames)
mydat <- readRDS("temp_results/lon_lat_ntest.rds")

lon  <- mydat$lon
lat <- mydat$lat
testYrsObs <- mydat$testYrsObs

nlon = length(lon)
nlat <- length(lat)
ntest <- length(testYrsObs)

testcat <- readRDS("temp_results/testcat.rds")
tblw <- testcat$tblw
tnrm <- testcat$tnrm
tabv <- testcat$tabv

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

# 2. Load the pieces
for (ir in 1:nprdctrs) {
  for (im in 1:nmodels) {
    path <- sprintf("temp_results/data_irun%02d_mod%02d.rds", ir, im)
    if (file.exists(path)) {
      dat <- readRDS(path)
      hhss_all[,,ir,im] <- dat$hhss
      test_probs_all[,,,,ir,im] <- dat$test_probs
      hsharpness_all[,,ir,im] <- dat$hsharpness
    }
  }
}

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]
      }
    }
  }
}
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)
    }
  }
}

test_years <- testYrsObs
source("templateWriteBestModels.R")
source("templateFrequency.R")
source("templateRealTimeHSS.R")


