#!/usr/bin/env bash
# =============================================================================
# Offline probe pre-trainer on generation-time SAE features (frozen base model)
# =============================================================================
# Single GPU:
#   bash experiment/scripts/probe/run_train_probe_gen.sh
# Multi-GPU:
#   NGPUS=4 bash experiment/scripts/probe/run_train_probe_gen.sh
# Common overrides:
#   RELATION=bathroom_toilet PROBE_EPOCHS=10 MAX_NEW_TOKENS_TRAIN=128 \
#     OUTPUT_DIR=./probe_outputs SAE_CHECKPOINT=/path/to/sae.ckpt POOL=mean bash ...
# =============================================================================

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}"
wandb login --relogin "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"


CUDA_VISIBLE_DEVICES="${CUDA_VISIBLE_DEVICES:-5}"
export CUDA_VISIBLE_DEVICES

RELATION="${RELATION:-bathroom_toilet}"
OUTPUT_DIR="${OUTPUT_DIR:-${PROJECT_ROOT}/probe_outputs}"
ADV_CONFIG="${ADV_CONFIG:-${PROJECT_ROOT}/experiment/adv_config.json}"
RUN_ID="${RUN_ID:-$(date +%Y%m%d_%H%M%S)}"
export RUN_ID
PROBE_EPOCHS="${PROBE_EPOCHS:-10}"
MAX_NEW_TOKENS_TRAIN="${MAX_NEW_TOKENS_TRAIN:-128}"
GEN_DO_SAMPLE="${GEN_DO_SAMPLE:-false}"
SAE_CHECKPOINT="${SAE_CHECKPOINT:-/data/caotue/SAE_checkpoints/16_d_model/last.ckpt}"
PROBE_OUTPUT="${PROBE_OUTPUT:-/data/caotue/SAE_checkpoints/16_d_model/probes_bathroom}"
POOL="${POOL:-max}"

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:-1}"
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

mkdir -p "$OUTPUT_DIR"

echo "=========================================="
echo "Probe pre-trainer (gen-time SAE features)"
echo "=========================================="
echo "  Project:   $PROJECT_ROOT"
echo "  Config:    $ADV_CONFIG"
echo "  Relation:  $RELATION"
echo "  Output:    $OUTPUT_DIR"
echo "  Epochs:    $PROBE_EPOCHS"
echo "  MaxTok:    $MAX_NEW_TOKENS_TRAIN"
echo "  Pool:      $POOL"
echo "  GPUs:      $NGPUS (visible=$VISIBLE_GPUS)"
echo "=========================================="

EXTRA_ARGS=(--config "$ADV_CONFIG" --relation "$RELATION" --output_dir "$OUTPUT_DIR")
EXTRA_ARGS+=(--probe_epochs "$PROBE_EPOCHS")
EXTRA_ARGS+=(--pool "$POOL")
EXTRA_ARGS+=(--adv.max_new_tokens_train "$MAX_NEW_TOKENS_TRAIN")
EXTRA_ARGS+=(--adv.gen_do_sample "$GEN_DO_SAMPLE")
[ -n "$SAE_CHECKPOINT" ] && EXTRA_ARGS+=(--adv.sae_checkpoint "$SAE_CHECKPOINT")
[ -n "$PROBE_OUTPUT" ] && EXTRA_ARGS+=(--probe_output "$PROBE_OUTPUT")

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.train_probe_gen "${EXTRA_ARGS[@]}" "$@"
else
    "$PYTHON_BIN" -m experiment.training.train_probe_gen "${EXTRA_ARGS[@]}" "$@"
fi