#!/usr/bin/env bash
set -euo pipefail

# Run all non-wrapper evaluation scripts for every completed experiment under
# runs/.  The script is written for Linux servers with bash 4.2.

cd "$(dirname "${BASH_SOURCE[0]}")"

PYTHON_BIN="${PYTHON_BIN:-python}"
DEVICE="${DEVICE:-cuda}"
EVAL_SPLIT="${EVAL_SPLIT:-test}"
NUM_WORKERS="${NUM_WORKERS:-4}"
CPU_REDUCE_WORKERS="${CPU_REDUCE_WORKERS:-}"
NUM_WORKERS_AUC="${NUM_WORKERS_AUC:-}"
BATCH_SIZE="${BATCH_SIZE:-}"
DATASET_SUBSET_SIZE="${DATASET_SUBSET_SIZE:-}"
DRY_RUN="${DRY_RUN:-0}"

# These attribution jobs can be expensive, but they are part of the evaluation
# surface in this repository. Set either variable to 0 to leave that family out.
RUN_EXTRA_INFO_ATTRIBUTION="${RUN_EXTRA_INFO_ATTRIBUTION:-1}"
RUN_SINGLE_DISEASE_MORTALITY_ATTRIBUTION="${RUN_SINGLE_DISEASE_MORTALITY_ATTRIBUTION:-1}"

common_args_base() {
  printf '%s\n' --run_path "$1" --eval_split "${EVAL_SPLIT}" --num_workers "${NUM_WORKERS}"
  if [[ -n "${BATCH_SIZE}" ]]; then
    printf '%s\n' --batch_size "${BATCH_SIZE}"
  fi
  if [[ -n "${DATASET_SUBSET_SIZE}" ]]; then
    printf '%s\n' --dataset_subset_size "${DATASET_SUBSET_SIZE}"
  fi
}

common_args_with_device() {
  common_args_base "$1"
  printf '%s\n' --device "${DEVICE}"
}

auc_args() {
  if [[ -n "${NUM_WORKERS_AUC}" ]]; then
    printf '%s\n' --num_workers_auc "${NUM_WORKERS_AUC}"
  fi
}

cpu_reduce_args() {
  if [[ -n "${CPU_REDUCE_WORKERS}" ]]; then
    printf '%s\n' --cpu_reduce_workers "${CPU_REDUCE_WORKERS}"
  fi
}

has_completed_dir() {
  local dir="$1"
  shift
  [[ -d "${dir}" ]] || return 1
  local required
  for required in "$@"; do
    [[ -s "${dir}/${required}" ]] || return 1
  done
}

run_command() {
  echo "    run: $*"
  if [[ "${DRY_RUN}" == "1" ]]; then
    return 0
  fi
  "$@"
}

run_dir_result_if_missing() {
  local label="$1"
  local result_dir="$2"
  local required_1="$3"
  local required_2="$4"
  shift 4

  if has_completed_dir "${result_dir}" "${required_1}" "${required_2}"; then
    echo "    skip ${label}: found ${result_dir}"
    return 0
  fi

  run_command "$@"
}

run_has_extra_info() {
  "${PYTHON_BIN}" - "$1" <<'PY'
import json
import sys
from pathlib import Path

cfg_path = Path(sys.argv[1]) / "train_config.json"
try:
    cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
except Exception:
    raise SystemExit(1)

extra = cfg.get("extra_info_types", [])
raise SystemExit(0 if isinstance(extra, list) and len(extra) > 0 else 1)
PY
}

run_is_all_future() {
  "${PYTHON_BIN}" - "$1" <<'PY'
import json
import sys
from pathlib import Path

cfg_path = Path(sys.argv[1]) / "train_config.json"
try:
    cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
except Exception:
    raise SystemExit(1)

mode = str(cfg.get("model_target_mode", "next_token")).lower()
raise SystemExit(0 if mode == "all_future" else 1)
PY
}

for run_path in runs/*; do
  [[ -d "${run_path}" ]] || continue

  echo "==> ${run_path}"
  if [[ ! -f "${run_path}/train_config.json" ]]; then
    echo "    skip run: missing train_config.json"
    continue
  fi
  if [[ ! -s "${run_path}/best_model.pt" ]]; then
    echo "    skip run: missing best_model.pt"
    continue
  fi

  common=()
  while IFS= read -r arg; do common+=("${arg}"); done < <(common_args_with_device "${run_path}")

  auc_extra=()
  while IFS= read -r arg; do auc_extra+=("${arg}"); done < <(auc_args)

  cpu_reduce_extra=()
  while IFS= read -r arg; do cpu_reduce_extra+=("${arg}"); done < <(cpu_reduce_args)

  run_dir_result_if_missing \
    "evaluate_auc.py" \
    "${run_path}" \
    "df_both.csv" \
    "df_auc_unpooled.csv" \
    "${PYTHON_BIN}" evaluate_auc.py "${common[@]}" "${auc_extra[@]}"

  run_dir_result_if_missing \
    "evaluate_auc_v2.py" \
    "${run_path}" \
    "df_auc_landmark.csv" \
    "df_auc_landmark_unpooled.csv" \
    "${PYTHON_BIN}" evaluate_auc_v2.py "${common[@]}" "${auc_extra[@]}"

  if ! run_is_all_future "${run_path}"; then
    echo "    skip attribution evaluations: model_target_mode is not all_future"
    continue
  fi

  if [[ "${RUN_EXTRA_INFO_ATTRIBUTION}" == "1" ]]; then
    if run_has_extra_info "${run_path}"; then
      run_dir_result_if_missing \
        "evaluate_extra_info_attribution.py" \
        "${run_path}/extra_info_attribution_${EVAL_SPLIT}" \
        "manifest.json" \
        "summary_extra_info_disease_parameters.csv" \
        "${PYTHON_BIN}" evaluate_extra_info_attribution.py "${common[@]}" "${cpu_reduce_extra[@]}"
    else
      echo "    skip evaluate_extra_info_attribution.py: run has no extra-info types"
    fi
  fi

  if [[ "${RUN_SINGLE_DISEASE_MORTALITY_ATTRIBUTION}" == "1" ]]; then
    run_dir_result_if_missing \
      "evaluate_single_disease_mortality_attribution.py" \
      "${run_path}/single_disease_mortality_parameters_${EVAL_SPLIT}_all_diseases" \
      "manifest.json" \
      "summary_by_disease_age_sex.csv" \
      "${PYTHON_BIN}" evaluate_single_disease_mortality_attribution.py "${common[@]}"
  fi
done

echo "All missing evaluations are complete."