#!/usr/bin/env bash
#
# Re-run every experiment invalidated by removing the assessment event token,
# requiring train-split RobustScale, and deleting the mixed distribution.
#
# The workflow runs:
#   1. seed-42 full factorial experiments;
#   2. seed-43/44 key-model experiments;
#   3. seed-42/43/44 disease-history T/O/S ablation;
#   4. seed-42/43/44 assessment + smoking + alcohol experiment;
#   5. AUC and calibration evaluation for all four campaigns.

set -euo pipefail

SCRIPT_DIR="${BASH_SOURCE[0]%/*}"
if [[ "$SCRIPT_DIR" == "${BASH_SOURCE[0]}" ]]; then
    SCRIPT_DIR="."
fi
SCRIPT_DIR="$(cd -- "$SCRIPT_DIR" && pwd)"
cd "$SCRIPT_DIR"

NUM_GPUS=1
NUM_GPUS_SET=0
GPU_CSV=""
NUM_CPUS=16
PYTHON_BIN="${PYTHON_BIN:-python}"
BASH_BIN="${BASH:-bash}"
DRY_RUN=0

FULL_FACTORIAL_CAMPAIGN="full_factorial_smoking_alcohol_bmi"
KEY_MODELS_CAMPAIGN="key_models_multiseed_smoking_alcohol_bmi"
HISTORY_CAMPAIGN="disease_history_ablation_no_extra"
EXTRA_INFO_CAMPAIGN="extra_info_assessment_smoking_alcohol_robust_multiseed"

usage() {
    cat <<'EOF'
Usage:
  bash rerun_all_required_experiments_linux.sh [options]

Options:
  --num-gpus N       Number of GPUs; uses GPU ids 0 through N-1 (default: 1).
  --gpus LIST        Explicit comma-separated GPU ids, for example 0,2.
                     Cannot be combined with --num-gpus.
  --num-cpus N       Total CPU workers shared across active GPUs (default: 16).
  --python PATH      Python executable used by every child script.
  --dry-run          Print commands without training or evaluation.
  -h, --help         Show this help message.

Examples:
  bash rerun_all_required_experiments_linux.sh --num-gpus 1 --num-cpus 16
  bash rerun_all_required_experiments_linux.sh --num-gpus 4 --num-cpus 32
  bash rerun_all_required_experiments_linux.sh --gpus 0,2 --num-cpus 16

CPU allocation:
  The script divides --num-cpus evenly across simultaneously active GPUs.
  For example, 4 GPUs and 32 CPUs gives each GPU job 8 workers.
EOF
}

while (($# > 0)); do
    case "$1" in
        --num-gpus)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --num-gpus requires a value." >&2
                exit 2
            }
            [[ -z "$GPU_CSV" ]] || {
                echo "ERROR: --num-gpus cannot be combined with --gpus." >&2
                exit 2
            }
            NUM_GPUS="$2"
            NUM_GPUS_SET=1
            shift 2
            ;;
        --gpus)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --gpus requires a value." >&2
                exit 2
            }
            ((NUM_GPUS_SET == 0)) || {
                echo "ERROR: --gpus cannot be combined with --num-gpus." >&2
                exit 2
            }
            GPU_CSV="$2"
            shift 2
            ;;
        --num-cpus)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --num-cpus requires a value." >&2
                exit 2
            }
            NUM_CPUS="$2"
            shift 2
            ;;
        --python)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --python requires a value." >&2
                exit 2
            }
            PYTHON_BIN="$2"
            shift 2
            ;;
        --dry-run)
            DRY_RUN=1
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "ERROR: unknown argument: $1" >&2
            usage >&2
            exit 2
            ;;
    esac
done

[[ "$NUM_GPUS" =~ ^[1-9][0-9]*$ ]] || {
    echo "ERROR: --num-gpus must be a positive integer." >&2
    exit 2
}
[[ "$NUM_CPUS" =~ ^[1-9][0-9]*$ ]] || {
    echo "ERROR: --num-cpus must be a positive integer." >&2
    exit 2
}

if [[ -z "$GPU_CSV" ]]; then
    for ((gpu_id = 0; gpu_id < NUM_GPUS; gpu_id++)); do
        if [[ -n "$GPU_CSV" ]]; then
            GPU_CSV+=","
        fi
        GPU_CSV+="$gpu_id"
    done
else
    IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
    NUM_GPUS="${#GPU_IDS[@]}"
    ((NUM_GPUS > 0)) || {
        echo "ERROR: --gpus must not be empty." >&2
        exit 2
    }
    for gpu_id in "${GPU_IDS[@]}"; do
        [[ "$gpu_id" =~ ^[0-9]+$ ]] || {
            echo "ERROR: invalid GPU id: $gpu_id" >&2
            exit 2
        }
    done
fi

((NUM_CPUS >= NUM_GPUS)) || {
    echo "ERROR: --num-cpus must be at least the number of GPUs." >&2
    exit 2
}
WORKERS_PER_GPU=$((NUM_CPUS / NUM_GPUS))

required_scripts=(
    train_batch_linux.sh
    train_key_models_multiseed_linux.sh
    train_disease_history_ablation_linux.sh
    train_extra_info_assessment_all_multiseed_linux.sh
    evaluate_all_runs_linux.sh
    evaluate_calibration_all_runs_linux.sh
)
for script in "${required_scripts[@]}"; do
    [[ -f "$SCRIPT_DIR/$script" ]] || {
        echo "ERROR: missing required script: $SCRIPT_DIR/$script" >&2
        exit 2
    }
done

run_command() {
    printf '>>'
    printf ' %q' "$@"
    printf '\n'
    if ((DRY_RUN == 0)); then
        "$@"
    fi
}

echo "GPUs: $GPU_CSV ($NUM_GPUS total)"
echo "CPUs: $NUM_CPUS total; $WORKERS_PER_GPU workers per GPU job"
echo "Python: $PYTHON_BIN"

run_command "$BASH_BIN" "$SCRIPT_DIR/train_batch_linux.sh" \
    --gpus "$GPU_CSV" \
    --seed 42 \
    --num-workers "$WORKERS_PER_GPU" \
    --python "$PYTHON_BIN"

run_command "$BASH_BIN" "$SCRIPT_DIR/train_key_models_multiseed_linux.sh" \
    --gpus "$GPU_CSV" \
    --seeds 43,44 \
    --num-workers "$WORKERS_PER_GPU" \
    --python "$PYTHON_BIN"

run_command "$BASH_BIN" "$SCRIPT_DIR/train_disease_history_ablation_linux.sh" \
    --gpus "$GPU_CSV" \
    --seeds 42,43,44 \
    --num-workers "$WORKERS_PER_GPU" \
    --python "$PYTHON_BIN"

run_command "$BASH_BIN" "$SCRIPT_DIR/train_extra_info_assessment_all_multiseed_linux.sh" \
    --gpus "$GPU_CSV" \
    --seeds 42,43,44 \
    --num-workers "$WORKERS_PER_GPU" \
    --python "$PYTHON_BIN"

campaigns=(
    "$FULL_FACTORIAL_CAMPAIGN"
    "$KEY_MODELS_CAMPAIGN"
    "$HISTORY_CAMPAIGN"
    "$EXTRA_INFO_CAMPAIGN"
)

for campaign in "${campaigns[@]}"; do
    runs_root="$SCRIPT_DIR/runs/$campaign"
    run_command "$BASH_BIN" "$SCRIPT_DIR/evaluate_all_runs_linux.sh" \
        --gpus "$GPU_CSV" \
        --runs-root "$runs_root" \
        --log-root "$SCRIPT_DIR/batch_logs/evaluate_all_required/$campaign" \
        --python "$PYTHON_BIN" \
        --num-workers "$WORKERS_PER_GPU" \
        --num-workers-auc "$WORKERS_PER_GPU"

    run_command "$BASH_BIN" "$SCRIPT_DIR/evaluate_calibration_all_runs_linux.sh" \
        --gpus "$GPU_CSV" \
        --runs-root "$runs_root" \
        --log-root "$SCRIPT_DIR/batch_logs/evaluate_calibration_all_required/$campaign" \
        --python "$PYTHON_BIN" \
        --num-workers "$WORKERS_PER_GPU" \
        --num-workers-calibration "$WORKERS_PER_GPU"
done

echo "All required re-training and evaluation workflows completed."