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

# Linux bash 5.2+ training-only script.
#
# Based on the existing runs, the objective/time/death-distribution checks are
# already covered. The remaining gap for the current proof chain is the
# extra-info ablation under the final candidate model:
#
#   all_future + relative time + mixed death/risk head
#
# This script only launches those missing training jobs. It intentionally does
# not call evaluate_*.py and does not add extra random seeds.

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

PYTHON_BIN="${PYTHON_BIN:-python}"
DEVICE="${DEVICE:-cuda}"
NUM_WORKERS="${NUM_WORKERS:-4}"
PROGRESS_INTERVAL="${PROGRESS_INTERVAL:-20}"

TIME_MODE="relative"
DIST_MODE="mixed"
SEED="42"
VALIDATION_QUERY_SEED="42"

COMMON_ARGS=(
  --data_prefix ukb
  --labels_file labels.csv
  --seed "${SEED}"
  --validation_query_seed "${VALIDATION_QUERY_SEED}"
  --train_eid_file ukb_train_eid.csv
  --val_eid_file ukb_val_eid.csv
  --test_eid_file ukb_test_eid.csv
  --min_history_events 1
  --min_future_events 1
  --n_embd 120
  --n_head 10
  --n_hist_layer 12
  --n_tab_layer 4
  --n_bins 16
  --extra_pool_reduce mean
  --dropout 0.0
  --batch_size 256
  --base_lr 0.0003
  --weight_decay 0.1
  --betas 0.9 0.99
  --grad_clip 1.0
  --max_epochs 200
  --warmup_epochs 10
  --patience 15
  --min_lr_ratio 0.1
  --num_workers "${NUM_WORKERS}"
  --device "${DEVICE}"
  --progress_interval "${PROGRESS_INTERVAL}"
)

already_trained() {
  local extra_file="$1"
  "${PYTHON_BIN}" - "$TIME_MODE" "$DIST_MODE" "$extra_file" "$SEED" "$VALIDATION_QUERY_SEED" <<'PY'
import json
import sys
from pathlib import Path

time_mode, dist_mode, extra_file, seed, validation_query_seed = sys.argv[1:6]
extra_name = Path(extra_file).name

for config_path in Path("runs").glob("*/train_config.json"):
    try:
        cfg = json.loads(config_path.read_text(encoding="utf-8"))
    except Exception:
        continue

    observed_query_seed = cfg.get(
        "all_future_validation_query_seed",
        cfg.get("validation_query_seed", -1),
    )

    if (
        cfg.get("model_target_mode") == "all_future"
        and cfg.get("time_mode") == time_mode
        and cfg.get("dist_mode") == dist_mode
        and Path(str(cfg.get("extra_info_types_file", ""))).name == extra_name
        and int(cfg.get("seed", -1)) == int(seed)
        and int(observed_query_seed) == int(validation_query_seed)
    ):
        print(config_path.parent)
        raise SystemExit(0)

raise SystemExit(1)
PY
}

train_if_missing() {
  local label="$1"
  local extra_file="$2"

  if [[ ! -f "${extra_file}" ]]; then
    echo "ERROR: missing extra-info type file: ${extra_file}" >&2
    return 2
  fi

  echo "==> Checking ${label}: ${TIME_MODE} ${DIST_MODE} all_future with ${extra_file}"
  if existing_run="$(already_trained "$extra_file")"; then
    echo "    skip: already trained at ${existing_run}"
    return 0
  fi

  echo "    train: ${label}"
  "${PYTHON_BIN}" train_all_future.py \
    "${COMMON_ARGS[@]}" \
    --time_mode "${TIME_MODE}" \
    --dist_mode "${DIST_MODE}" \
    --extra_info_types_file "${extra_file}"
}

# Already present in runs/:
# - next-token objective checks under SAB, plus older absolute extra ablations
# - all-future absolute/relative x exponential/weibull/mixed under SAB
#
# Still needed:
# - final all-future relative+mixed extra-info ablations beyond the existing
#   SAB baseline. These close the disease-only question without expanding seed
#   count or running downstream evaluation.
train_if_missing "true_disease_only" "extra_info_types_none.txt"
train_if_missing "assessment_only_extra" "extra_info_types_assessment_only.txt"
train_if_missing "exposure_only_extra" "extra_info_types_exposure_only.txt"
train_if_missing "all_extra_info" "extra_info_types_all.txt"

echo "All requested training-only missing configurations are done."