args <- commandArgs(trailingOnly = FALSE)
script_arg <- grep("^--file=", args, value = TRUE)
if (length(script_arg) > 0) {
  script_path <- normalizePath(sub("^--file=", "", script_arg[1]))
  setwd(normalizePath(file.path(dirname(script_path), "..")))
}

library(reticulate)
use_python(Sys.getenv("RETICULATE_PYTHON"), required = TRUE)

tf <- import("tensorflow", convert = FALSE)
helper <- import_from_path("tf_savedmodel_helper", path = file.path(getwd(), "scripts"), convert = TRUE)

mcc_score <- function(pred, real) {
  tp <- as.numeric(sum(pred == 1 & real == 1))
  tn <- as.numeric(sum(pred == 0 & real == 0))
  fp <- as.numeric(sum(pred == 1 & real == 0))
  fn <- as.numeric(sum(pred == 0 & real == 1))
  denom <- sqrt((tp + fp) * (tp + fn) * (tn + fp) * (tn + fn))
  if (is.na(denom) || denom == 0) return(NA_real_)
  (tp * tn - fp * fn) / denom
}

get_input_name <- function(model_dir) {
  model <- tf$saved_model$load(model_dir)
  serving <- model$signatures$get("serving_default")
  sig_text <- py_str(serving$structured_input_signature)
  out_text <- py_str(serving$structured_outputs)

  cat("\nModel:", model_dir, "\n")
  cat("Input signature:", sig_text, "\n")
  cat("Output signature:", out_text, "\n")

  input_name <- sub(".*'([^']+)': TensorSpec.*", "\\1", sig_text)
  cat("Input name:", input_name, "\n")
  input_name
}

run_eval <- function(model_dir, x_file, y_file, label) {
  x <- readRDS(x_file)
  y <- readRDS(y_file)

  input_name <- get_input_name(model_dir)
  pred_prob <- helper$predict_saved_model(model_dir, input_name, x)

  y_real <- max.col(y) - 1
  y_pred <- max.col(pred_prob) - 1

  acc <- mean(y_real == y_pred)
  mcc <- mcc_score(y_pred, y_real)

  cat("\n====", label, "====\n")
  cat("n_test:", length(y_real), "\n")
  cat("accuracy:", round(acc, 4), "\n")
  cat("mcc:", round(mcc, 4), "\n")
  print(table(real = y_real, pred = y_pred))

  invisible(list(prob = pred_prob, real = y_real, pred = y_pred))
}

c1_res <- run_eval("weight/CNN/model_c1", "model/CNN/c1_test.RDS", "model/CNN/c1_test_y.RDS", "CTLA-4")
p1_res <- run_eval("weight/CNN/model_p1", "model/CNN/p1_test.RDS", "model/CNN/p1_test_y.RDS", "PD-1")

saveRDS(c1_res, "model/CNN/c1_tf218_inference_result.RDS")
saveRDS(p1_res, "model/CNN/p1_tf218_inference_result.RDS")

cat("\nCNN inference OK\n")