#!/bin/bash
#SBATCH --job-name=visualize_sae_features
#SBATCH --output=./log_slurm/result/visualize_sae_features.txt
#SBATCH --error=./log_slurm/error/visualize_sae_features.txt
#SBATCH --ntasks=1
#SBATCH --gpus=1
#SBATCH --nodes=1
#SBATCH --cpus-per-task=20

# =============================================================================
# Visualize SAE Features (LLaVA, multi-GPU)
#
# Runs training/visualize_multilayer_sae_features.py to produce HTML reports
# for each specified feature.
#
# DATA_MODE choices:
#   toilet  (default) — pbcong/bathroom-toilet positives + CC3M negatives.
#                       Needs IMAGE_FOLDER. Set HF_DATASET + SPLIT if you want
#                       CC3M captions for negatives in caption mode.
#   cc3m              — full CC3M via HF_DATASET + LOCAL_VAL_PATH.
#   coco              — COCO via HF_DATASET + LOCAL_VAL_PATH.
#   folder            — plain image folder via DATA_DIR.
#
# CAPTION_MODE choices:
#   generated  (default) — model generates caption; teacher-forced re-pass
#                          collects residual activations.
#   caption              — use stored caption (CC3M txt / COCO sentences /
#                          pbcong/bathroom-toilet caption field).
#
# Usage:
#   bash training/scripts/visualize_multilayer_sae_features.sh
#
# Override any variable inline:
#   DATA_MODE=cc3m HF_DATASET=pixparse/cc3m-wds \
#   LOCAL_VAL_PATH=/path/to/cc3m/train \
#   FEATURE_IDS="0 1 42" \
#   bash training/scripts/visualize_multilayer_sae_features.sh
# =============================================================================

# ── GPU ──────────────────────────────────────────────────────────────────────
NUM_GPUS="${NUM_GPUS:-1}"
DEVICE_ID="${DEVICE_ID:-0}"

# ── Model & SAE ──────────────────────────────────────────────────────────────
SAE_CKPT="${SAE_CKPT:-training/multilayer_sae_ckpt/last.ckpt}"
MODEL_NAME="${MODEL_NAME:-llava-hf/llava-1.5-7b-hf}"
DTYPE="${DTYPE:-float16}"

# ── Data mode ────────────────────────────────────────────────────────────────
DATA_MODE="${DATA_MODE:-toilet}"        # toilet | cc3m | coco | folder
CAPTION_MODE="${CAPTION_MODE:-generated}"  # generated | caption
NUM_WORKERS="${NUM_WORKERS:-8}"

# ── Data — toilet ────────────────────────────────────────────────────────────
IMAGE_FOLDER="${IMAGE_FOLDER:-CC3M-Dataset/cc3m_images/train}"
TOILET_MODE="${TOILET_MODE:-toilet}"    # toilet | bathroom
NUM_NEGATIVES="${NUM_NEGATIVES:-10000}"

# ── Data — CC3M / COCO (also used for toilet-negative captions) ──────────────
HF_DATASET="${HF_DATASET:-}"           # e.g. pixparse/cc3m-wds or yerevann/coco-karpathy
LOCAL_VAL_PATH="${LOCAL_VAL_PATH:-}"   # local image root for HF dataset
DATA_DIR="${DATA_DIR:-}"               # [folder mode] plain image directory
SPLIT="${SPLIT:-train}"

# ── Hook point (REQUIRED) ────────────────────────────────────────────────────
HOOK_POINT="${HOOK_POINT:-model.language_model.layers.19.hook_resid_post}"

# ── Feature selection (REQUIRED) ────────────────────────────────────────────
FEATURE_IDS="${FEATURE_IDS:-49199 18633 6307 52414 879 29368 55162 41137 39220 8437}"

# ── Processing ───────────────────────────────────────────────────────────────
BATCH_SIZE="${BATCH_SIZE:-4}"
SAE_BATCH="${SAE_BATCH:-4096}"
THRESHOLD="${THRESHOLD:-1e-3}"
MAX_BATCHES="${MAX_BATCHES:-}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-128}"

# ── Visualisation ────────────────────────────────────────────────────────────
OUTPUT_DIR="${OUTPUT_DIR:-training/visualize_features}"
TOP_IMAGES="${TOP_IMAGES:-20}"
TOP_TEXTS="${TOP_TEXTS:-20}"
BUFFER="${BUFFER:-10}"

# =============================================================================
# Validation
# =============================================================================
cd "$(dirname "$0")/../.."

if [ ! -f "${SAE_CKPT}" ]; then
    echo "Error: SAE checkpoint not found: ${SAE_CKPT}" >&2
    exit 1
fi

case "${DATA_MODE}" in
    toilet)
        if [ ! -d "${IMAGE_FOLDER}" ]; then
            echo "Error: IMAGE_FOLDER not found: ${IMAGE_FOLDER}" >&2
            exit 1
        fi
        ;;
    cc3m|coco)
        if [ -z "${HF_DATASET}" ]; then
            echo "Error: DATA_MODE=${DATA_MODE} requires HF_DATASET." >&2
            exit 1
        fi
        ;;
    folder)
        if [ ! -d "${DATA_DIR}" ]; then
            echo "Error: DATA_DIR not found: ${DATA_DIR}" >&2
            exit 1
        fi
        ;;
    *)
        echo "Error: Unknown DATA_MODE=${DATA_MODE}. Choose: toilet | cc3m | coco | folder" >&2
        exit 1
        ;;
esac

if [ -z "${HOOK_POINT}" ]; then
    echo "Error: HOOK_POINT must be set." >&2
    exit 1
fi

if [ -z "${FEATURE_IDS}" ]; then
    echo "Error: FEATURE_IDS must be set." >&2
    exit 1
fi

# =============================================================================
# Environment
# =============================================================================
export HF_HOME="${HF_HOME:-${HOME}/scratch/hf_home}"
export PYTHONPATH="$(pwd):${PYTHONPATH:-}"

if [ -f .env ]; then
    set -a; source .env; set +a
fi

# =============================================================================
# Build argument list
# =============================================================================
ARGS=(
    --sae_ckpt        "${SAE_CKPT}"
    --model_name      "${MODEL_NAME}"
    --device_id       "${DEVICE_ID}"
    --dtype           "${DTYPE}"
    --data_mode       "${DATA_MODE}"
    --caption_mode    "${CAPTION_MODE}"
    --num_workers     "${NUM_WORKERS}"
    --hook_point      "${HOOK_POINT}"
    --output_dir      "${OUTPUT_DIR}"
    --feature_ids     ${FEATURE_IDS}
    --batch_size      "${BATCH_SIZE}"
    --sae_batch       "${SAE_BATCH}"
    --threshold       "${THRESHOLD}"
    --max_new_tokens  "${MAX_NEW_TOKENS}"
    --top_images      "${TOP_IMAGES}"
    --top_texts       "${TOP_TEXTS}"
    --buffer          "${BUFFER}"
    --split           "${SPLIT}"
)

# ── Data-mode-specific args ──────────────────────────────────────────────────
case "${DATA_MODE}" in
    toilet)
        ARGS+=(
            --image_folder  "${IMAGE_FOLDER}"
            --toilet_mode   "${TOILET_MODE}"
            --num_negatives "${NUM_NEGATIVES}"
        )
        # Optionally pass CC3M HF dataset for negative caption lookup
        if [ -n "${HF_DATASET}" ]; then
            ARGS+=(--hf_dataset "${HF_DATASET}")
        fi
        ;;
    cc3m|coco)
        ARGS+=(--hf_dataset "${HF_DATASET}")
        if [ -n "${LOCAL_VAL_PATH}" ]; then
            ARGS+=(--local_val_path "${LOCAL_VAL_PATH}")
        fi
        ;;
    folder)
        ARGS+=(--data_dir "${DATA_DIR}")
        ;;
esac

if [ -n "${MAX_BATCHES}" ]; then
    ARGS+=(--max_batches "${MAX_BATCHES}")
fi

# =============================================================================
# Run
# =============================================================================
if [ "${NUM_GPUS}" -gt 1 ]; then
    echo "Launching with torchrun on ${NUM_GPUS} GPUs..."
    torchrun --nproc_per_node="${NUM_GPUS}" -m training.visualize_multilayer_sae_features "${ARGS[@]}"
else
    echo "Launching single-GPU mode..."
    python -m training.visualize_multilayer_sae_features "${ARGS[@]}"
fi