#!/usr/bin/env bash
# =============================================================================
# Eval a trained gen-time probe on a held-out split (default: val)
# =============================================================================
# Usage:
#   PROBE_CHECKPOINT=/path/to/probes_gen_<rel>.pt \
#     bash experiment/scripts/probe/run_eval_probe_gen.sh
# Common overrides:
#   RELATION=bathroom_toilet SPLIT=val NGPUS=4 \
#     PROBE_CHECKPOINT=... OUTPUT_JSON=eval.json \
#     bash experiment/scripts/probe/run_eval_probe_gen.sh
# =============================================================================

set -euo pipefail

PROJECT_ROOT="$(cd "$(dirname "$0")/../../.." && pwd)"
export PYTHONPATH="${PROJECT_ROOT}:${PYTHONPATH:-}"
export TOKENIZERS_PARALLELISM="${TOKENIZERS_PARALLELISM:-false}"
export TORCHINDUCTOR_CACHE_DIR="${TORCHINDUCTOR_CACHE_DIR:-${HOME}/scratch/.cache/torchinductor}"
export TRITON_CACHE_DIR="${TRITON_CACHE_DIR:-${HOME}/scratch/.cache/triton}"

RELATION="${RELATION:-bathroom_toilet}"
ADV_CONFIG="${ADV_CONFIG:-${PROJECT_ROOT}/experiment/adv_config.json}"
SPLIT="${SPLIT:-val}"
PROBE_CHECKPOINT="${PROBE_CHECKPOINT:-}"
SAE_CHECKPOINT="${SAE_CHECKPOINT:-}"
OUTPUT_JSON="${OUTPUT_JSON:-}"
MAX_EVAL_SAMPLES="${MAX_EVAL_SAMPLES:-}"
MAX_NEW_TOKENS_TRAIN="${MAX_NEW_TOKENS_TRAIN:-}"
GEN_DO_SAMPLE="${GEN_DO_SAMPLE:-}"
POOL="${POOL:-max}"

if [ -z "$PROBE_CHECKPOINT" ]; then
    echo "ERROR: PROBE_CHECKPOINT is required (path to probes_gen_<rel>.pt)" >&2
    exit 2
fi
if [ ! -f "$PROBE_CHECKPOINT" ]; then
    echo "ERROR: probe checkpoint not found: $PROBE_CHECKPOINT" >&2
    exit 2
fi

PYTHON_BIN="${PYTHON_BIN:-${PROJECT_ROOT}/.venv/bin/python}"
TORCHRUN_BIN="${TORCHRUN_BIN:-${PROJECT_ROOT}/.venv/bin/torchrun}"
[ -x "$PYTHON_BIN" ] || PYTHON_BIN="python3"
[ -x "$TORCHRUN_BIN" ] || TORCHRUN_BIN="torchrun"

count_visible_gpus() {
    local n
    n=$(nvidia-smi -L 2>/dev/null | grep -c '^GPU' || true)
    n=$(echo "$n" | tr -d '[:space:]')
    [ -z "$n" ] && n=1
    [ "$n" -lt 1 ] && n=1
    echo "$n"
}

VISIBLE_GPUS=$(count_visible_gpus)
NGPUS_RAW="${NGPUS:-auto}"
if [ -z "$NGPUS_RAW" ] || [ "$NGPUS_RAW" = "auto" ]; then
    NGPUS=$VISIBLE_GPUS
else
    NGPUS=$NGPUS_RAW
fi
[ "$NGPUS" -lt 1 ] && NGPUS=1
[ "$NGPUS" -gt "$VISIBLE_GPUS" ] && NGPUS=$VISIBLE_GPUS

echo "=========================================="
echo "Probe eval (gen-time SAE features)"
echo "=========================================="
echo "  Project:   $PROJECT_ROOT"
echo "  Config:    $ADV_CONFIG"
echo "  Relation:  $RELATION"
echo "  Split:     $SPLIT"
echo "  ProbeCkpt: $PROBE_CHECKPOINT"
echo "  Pool:      $POOL"
echo "  GPUs:      $NGPUS (visible=$VISIBLE_GPUS)"
echo "=========================================="

EXTRA_ARGS=(--config "$ADV_CONFIG" --relation "$RELATION")
EXTRA_ARGS+=(--probe_checkpoint "$PROBE_CHECKPOINT" --split "$SPLIT" --pool "$POOL")
[ -n "$SAE_CHECKPOINT" ]      && EXTRA_ARGS+=(--adv.sae_checkpoint "$SAE_CHECKPOINT")
[ -n "$MAX_NEW_TOKENS_TRAIN" ] && EXTRA_ARGS+=(--adv.max_new_tokens_train "$MAX_NEW_TOKENS_TRAIN")
[ -n "$GEN_DO_SAMPLE" ]        && EXTRA_ARGS+=(--adv.gen_do_sample "$GEN_DO_SAMPLE")
[ -n "$MAX_EVAL_SAMPLES" ]     && EXTRA_ARGS+=(--max_eval_samples "$MAX_EVAL_SAMPLES")
[ -n "$OUTPUT_JSON" ]          && EXTRA_ARGS+=(--output_json "$OUTPUT_JSON")

if [ "$NGPUS" -gt 1 ]; then
    export MASTER_ADDR="${MASTER_ADDR:-127.0.0.1}"
    [ -z "${MASTER_PORT:-}" ] && export MASTER_PORT=$((29500 + RANDOM % 1000))
    echo "  MASTER_ADDR: $MASTER_ADDR  MASTER_PORT: $MASTER_PORT"
    "$TORCHRUN_BIN" --nproc_per_node="$NGPUS" \
        --master_addr="$MASTER_ADDR" --master_port="$MASTER_PORT" \
        -m experiment.training.eval_probe_gen "${EXTRA_ARGS[@]}" "$@"
else
    "$PYTHON_BIN" -m experiment.training.eval_probe_gen "${EXTRA_ARGS[@]}" "$@"
fi