#!/usr/bin/env bash
#
# Recursively evaluate calibration for every completed all_future run.
#
# A runnable run contains:
#   - train_config.json with model_target_mode="all_future"
#   - best_model.pt
#
# calibration_evaluation_summary.json is written last by the evaluator and is
# used as the completion marker. Existing non-empty markers are skipped unless
# --force is supplied. next_token/Delphi2M runs are intentionally skipped.
#
# Jobs assigned to the same GPU run sequentially; different GPUs run in
# parallel.
#
# Examples:
#   bash evaluate_calibration_all_runs_linux.sh --gpus 0
#   bash evaluate_calibration_all_runs_linux.sh --gpus 0,1,2,3
#   bash evaluate_calibration_all_runs_linux.sh --gpus 0,1 --dry-run
#

set -uo pipefail

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"

RUNS_ROOT="$SCRIPT_DIR/runs"
LOG_ROOT="$SCRIPT_DIR/batch_logs/evaluate_calibration_all_runs"
GPU_CSV="0"
PYTHON_BIN="${PYTHON_BIN:-python}"
NUM_WORKERS=4
NUM_WORKERS_CALIBRATION=0
BATCH_SIZE=128
DISEASE_CHUNK_SIZE=64
HORIZONS=""
USE_AMP=0
FORCE=0
DRY_RUN=0

COMPLETION_FILE="calibration_evaluation_summary.json"
EVALUATOR="$SCRIPT_DIR/evaluate_calibration.py"

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

Options:
  --gpus LIST                 Comma-separated GPU ids (default: 0).
  --runs-root PATH            Root directory scanned recursively
                              (default: ./runs).
  --log-root PATH             Evaluation log root.
  --python PATH               Python executable
                              (default: $PYTHON_BIN or python).
  --num-workers N             DataLoader workers per job (default: 4).
  --num-workers-calibration N CPU calibration workers per job. Default: 0,
                              which divides all logical CPUs across GPUs.
  --batch-size N              Evaluation batch size (default: 128).
  --disease-chunk-size N      Disease projection chunk size (default: 64).
  --horizons LIST             Optional comma-separated horizons in years.
  --use-amp                   Force CUDA automatic mixed precision.
  --force                     Recompute runs with an existing completion file.
  --dry-run                   Discover and print pending jobs only.
  -h, --help                  Show this help message.

Only all_future runs are evaluated. next_token runs are skipped.
Completion marker: calibration_evaluation_summary.json
EOF
}

while (($# > 0)); do
    case "$1" in
        --gpus)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --gpus requires a value." >&2
                exit 2
            }
            GPU_CSV="$2"
            shift 2
            ;;
        --runs-root)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --runs-root requires a value." >&2
                exit 2
            }
            RUNS_ROOT="$2"
            shift 2
            ;;
        --log-root)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --log-root requires a value." >&2
                exit 2
            }
            LOG_ROOT="$2"
            shift 2
            ;;
        --python)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --python requires a value." >&2
                exit 2
            }
            PYTHON_BIN="$2"
            shift 2
            ;;
        --num-workers)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --num-workers requires a value." >&2
                exit 2
            }
            NUM_WORKERS="$2"
            shift 2
            ;;
        --num-workers-calibration)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --num-workers-calibration requires a value." >&2
                exit 2
            }
            NUM_WORKERS_CALIBRATION="$2"
            shift 2
            ;;
        --batch-size)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --batch-size requires a value." >&2
                exit 2
            }
            BATCH_SIZE="$2"
            shift 2
            ;;
        --disease-chunk-size)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --disease-chunk-size requires a value." >&2
                exit 2
            }
            DISEASE_CHUNK_SIZE="$2"
            shift 2
            ;;
        --horizons)
            [[ $# -ge 2 ]] || {
                echo "ERROR: --horizons requires a value." >&2
                exit 2
            }
            HORIZONS="$2"
            shift 2
            ;;
        --use-amp)
            USE_AMP=1
            shift
            ;;
        --force)
            FORCE=1
            shift
            ;;
        --dry-run)
            DRY_RUN=1
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "ERROR: unknown argument: $1" >&2
            usage >&2
            exit 2
            ;;
    esac
done

[[ -n "$GPU_CSV" ]] || {
    echo "ERROR: --gpus must not be empty." >&2
    exit 2
}
[[ "$NUM_WORKERS" =~ ^[0-9]+$ ]] || {
    echo "ERROR: --num-workers must be a non-negative integer." >&2
    exit 2
}
[[ "$NUM_WORKERS_CALIBRATION" =~ ^[0-9]+$ ]] || {
    echo "ERROR: --num-workers-calibration must be a non-negative integer." >&2
    exit 2
}
[[ "$BATCH_SIZE" =~ ^[1-9][0-9]*$ ]] || {
    echo "ERROR: --batch-size must be a positive integer." >&2
    exit 2
}
[[ "$DISEASE_CHUNK_SIZE" =~ ^[1-9][0-9]*$ ]] || {
    echo "ERROR: --disease-chunk-size must be a positive integer." >&2
    exit 2
}
[[ -d "$RUNS_ROOT" ]] || {
    echo "ERROR: runs root does not exist: $RUNS_ROOT" >&2
    exit 2
}
[[ -f "$EVALUATOR" ]] || {
    echo "ERROR: missing evaluator: $EVALUATOR" >&2
    exit 2
}
command -v "$PYTHON_BIN" >/dev/null 2>&1 || {
    echo "ERROR: Python executable not found: $PYTHON_BIN" >&2
    exit 2
}

RUNS_ROOT="$(cd -- "$RUNS_ROOT" && pwd)"
if [[ "$LOG_ROOT" != /* ]]; then
    LOG_ROOT="$SCRIPT_DIR/$LOG_ROOT"
fi

IFS=',' read -r -a GPU_IDS <<< "$GPU_CSV"
declare -A SEEN_GPUS=()
for gpu in "${GPU_IDS[@]}"; do
    [[ -n "$gpu" && "$gpu" =~ ^[A-Za-z0-9._:-]+$ ]] || {
        echo "ERROR: invalid GPU id: $gpu" >&2
        exit 2
    }
    [[ -z "${SEEN_GPUS[$gpu]+x}" ]] || {
        echo "ERROR: duplicate GPU id: $gpu" >&2
        exit 2
    }
    SEEN_GPUS["$gpu"]=1
done

if ((NUM_WORKERS_CALIBRATION == 0)); then
    TOTAL_CPUS="$(nproc)"
    NUM_WORKERS_CALIBRATION=$(( (TOTAL_CPUS + ${#GPU_IDS[@]} - 1) / ${#GPU_IDS[@]} ))
    if ((NUM_WORKERS_CALIBRATION < 1)); then
        NUM_WORKERS_CALIBRATION=1
    fi
fi

declare -a JOB_RUN_DIRS=()
declare -a JOB_LOG_FILES=()

add_job() {
    local run_dir="$1"
    local relative_run
    if [[ "$run_dir" == "$RUNS_ROOT" ]]; then
        relative_run="_root"
    else
        relative_run="${run_dir#"$RUNS_ROOT"/}"
    fi
    JOB_RUN_DIRS+=("$run_dir")
    JOB_LOG_FILES+=("$LOG_ROOT/$relative_run/evaluate_calibration.log")
}

run_count=0
incomplete_count=0
next_token_count=0
invalid_config_count=0
existing_count=0

while IFS= read -r -d '' config_path; do
    run_dir="${config_path%/train_config.json}"
    ((run_count += 1))

    if [[ ! -f "$run_dir/best_model.pt" ]]; then
        echo "[SKIP] Incomplete run without best_model.pt: $run_dir"
        ((incomplete_count += 1))
        continue
    fi

    if ! target_mode="$(
        "$PYTHON_BIN" -c \
            'import json,sys; print(str(json.load(open(sys.argv[1], encoding="utf-8")).get("model_target_mode", "next_token")).lower())' \
            "$config_path"
    )"; then
        echo "[SKIP] Invalid train_config.json: $config_path" >&2
        ((invalid_config_count += 1))
        continue
    fi

    if [[ "$target_mode" != "all_future" ]]; then
        echo "[SKIP] model_target_mode=$target_mode: $run_dir"
        ((next_token_count += 1))
        continue
    fi

    if ((FORCE == 0)) && [[ -s "$run_dir/$COMPLETION_FILE" ]]; then
        ((existing_count += 1))
        continue
    fi

    add_job "$run_dir"
done < <(find "$RUNS_ROOT" -type f -name "train_config.json" -print0)

if ((DRY_RUN == 0)); then
    mkdir -p "$LOG_ROOT"
fi

print_command() {
    printf '%q ' "$@"
    printf '\n'
}

run_job() {
    local job_index="$1"
    local gpu="$2"
    local run_dir="${JOB_RUN_DIRS[$job_index]}"
    local log_file="${JOB_LOG_FILES[$job_index]}"
    local -a command=(
        "$PYTHON_BIN"
        -u
        "$EVALUATOR"
        --run_path "$run_dir"
        --output_path "$run_dir"
        --eval_split test
        --device cuda
        --num_workers "$NUM_WORKERS"
        --num_workers_calibration "$NUM_WORKERS_CALIBRATION"
        --batch_size "$BATCH_SIZE"
        --disease_chunk_size "$DISEASE_CHUNK_SIZE"
    )

    if [[ -n "$HORIZONS" ]]; then
        command+=(--horizons "$HORIZONS")
    fi
    if ((USE_AMP)); then
        command+=(--use_amp)
    fi
    if ((FORCE)); then
        command+=(--force)
    fi

    echo "[$(date '+%F %T')] START gpu=$gpu"
    echo "  run=$run_dir"

    if ((DRY_RUN)); then
        printf '  CUDA_VISIBLE_DEVICES=%q ' "$gpu"
        print_command "${command[@]}"
        return 0
    fi

    mkdir -p "$(dirname -- "$log_file")"
    if CUDA_VISIBLE_DEVICES="$gpu" PYTHONUNBUFFERED=1 \
        "${command[@]}" >"$log_file" 2>&1; then
        if [[ -s "$run_dir/$COMPLETION_FILE" ]]; then
            echo "[$(date '+%F %T')] DONE  gpu=$gpu"
            return 0
        fi
        echo "[$(date '+%F %T')] FAIL  gpu=$gpu" >&2
        echo "  Missing completion marker: $run_dir/$COMPLETION_FILE" >&2
        echo "  See: $log_file" >&2
        return 1
    else
        local exit_code=$?
        echo "[$(date '+%F %T')] FAIL  gpu=$gpu exit=$exit_code" >&2
        echo "  See: $log_file" >&2
        return "$exit_code"
    fi
}

worker() {
    local slot="$1"
    local gpu="${GPU_IDS[$slot]}"
    local job_index
    local failed=0

    for ((
        job_index = slot;
        job_index < ${#JOB_RUN_DIRS[@]};
        job_index += ${#GPU_IDS[@]}
    )); do
        run_job "$job_index" "$gpu" || failed=1
    done
    return "$failed"
}

echo "Runs root: $RUNS_ROOT"
echo "GPUs: ${GPU_IDS[*]}"
echo "Calibration CPU workers per GPU job: $NUM_WORKERS_CALIBRATION"
echo "Runs discovered: $run_count"
echo "Incomplete runs skipped: $incomplete_count"
echo "next_token runs skipped: $next_token_count"
echo "Invalid configs skipped: $invalid_config_count"
echo "Existing calibration results skipped: $existing_count"
echo "Pending all_future evaluations: ${#JOB_RUN_DIRS[@]}"
echo "Log root: $LOG_ROOT"
echo

if ((${#JOB_RUN_DIRS[@]} == 0)); then
    echo "No pending all_future calibration evaluations."
    exit 0
fi

declare -a WORKER_PIDS=()
for ((slot = 0; slot < ${#GPU_IDS[@]}; slot++)); do
    worker "$slot" &
    WORKER_PIDS+=("$!")
done

overall_status=0
for pid in "${WORKER_PIDS[@]}"; do
    wait "$pid" || overall_status=1
done

if ((overall_status != 0)); then
    echo "One or more calibration evaluations failed." >&2
    echo "Inspect logs under: $LOG_ROOT" >&2
    exit 1
fi

if ((DRY_RUN)); then
    echo "Dry run completed successfully."
else
    echo "All pending all_future calibration evaluations completed."
fi