#!/bin/bash

#PBS -N visualize_probe_features
#PBS -l select=1:ngpus=8
#PBS -l walltime=12:00:00
#PBS -q ic102
#PBS -P 71001002
#PBS -j oe

# =============================================================================
# PBS Job File: Visualize Probe Features — Multi-Layer (LLaVA, multi-GPU)
#
# Runs training/visualize_probe_features.py which:
#   1. Loads linear probe checkpoints for each layer → top-k features by weight
#   2. Runs a single forward pass per layer to collect top activations
#   3. Generates ONE self-contained interactive HTML:
#        Layers → Top Features (+ probe weight) → Image patches / Text tokens
#
# DATA MODES  (set DATA_MODE below)
# ---------
#   toilet  — only HF "pbcong/bathroom-toilet" images matching OBJECT_MODE.
#             Set IMAGE_FOLDER to the local CC3M image directory.
#             Supports CAPTION_MODE=generated|dataset.
#
#   cc3m    — full CC3M or COCO via HF dataset + local path, OR a plain image
#             folder.  Set HF_DATASET+LOCAL_VAL_PATH  -or-  DATA_DIR.
#
# To submit:
#   qsub training/scripts/run_visualize_probe_features.pbs
#
# Override any variable before submission:
#   DATA_MODE=cc3m LAYERS="0 1 2 3 4 5 6" \
#   qsub training/scripts/run_visualize_probe_features.pbs
# =============================================================================

cd ${PBS_O_WORKDIR}

# Load CUDA
module load cuda/12.6.2
export CUDA_HOME="${CUDA_HOME:-$(dirname $(dirname $(which nvcc 2>/dev/null)))}"

# Activate environment
source "${HOME}/scratch/testing/multilayer-sae/venv/bin/activate"

# Ensure paths
export PATH=$HOME/.local/bin:$PATH
export HF_HOME="${HF_HOME:-${HOME}/scratch/hf_home}"
export PYTHONPATH="${PBS_O_WORKDIR}:${PYTHONPATH:-}"

if [ -f "${PBS_O_WORKDIR}/.env" ]; then
    set -a
    source "${PBS_O_WORKDIR}/.env"
    set +a
fi

# Create log directory
mkdir -p logs/visualize

# =============================================================================
# Hyperparameters — override via env before qsub
# =============================================================================

# ── GPU ───────────────────────────────────────────────────────────────────────
export NUM_GPUS="${NUM_GPUS:-8}"

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

# ── Probe ─────────────────────────────────────────────────────────────────────
export PROBE_DIR="${PROBE_DIR:-training/multilayer_sae_ckpt}"
export PROBE_INPUT_DIM="${PROBE_INPUT_DIM:-65536}"
export LAYERS="${LAYERS:-0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31}"
export TOP_PROBE_K="${TOP_PROBE_K:-10}"

# ── Data mode: toilet | cc3m | coco | folder ─────────────────────────────────
export DATA_MODE="${DATA_MODE:-cc3m}"

# ── Data — toilet mode ────────────────────────────────────────────────────────
export IMAGE_FOLDER="${IMAGE_FOLDER:-/home/users/ntu/cong045/scratch/testing/hallucination/CC3M-Dataset/cc3m_images/train}"
export OBJECT_MODE="${OBJECT_MODE:-toilet}"          # toilet | bathroom | both
export CAPTION_MODE="${CAPTION_MODE:-dataset}"     # generated | dataset
export MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-128}"
export NUM_NEGATIVES="${NUM_NEGATIVES:-10000}"

# ── Data — cc3m mode (HF dataset + local path) ───────────────────────────────
export HF_DATASET="${HF_DATASET:-pixparse/cc3m-wds}"
export LOCAL_VAL_PATH="${LOCAL_VAL_PATH:-/home/users/ntu/cong045/scratch/testing/hallucination/CC3M-Dataset/cc3m_images/train}"
export SPLIT="${SPLIT:-train}"

# ── Data — cc3m mode (plain image folder, alternative to HF_DATASET) ─────────
export DATA_DIR="${DATA_DIR:-}"

# ── Common data ───────────────────────────────────────────────────────────────
export NUM_WORKERS="${NUM_WORKERS:-32}"

# ── Processing ────────────────────────────────────────────────────────────────
export BATCH_SIZE="${BATCH_SIZE:-256}"
export SAE_BATCH="${SAE_BATCH:-4096}"
export THRESHOLD="${THRESHOLD:-1e-3}"
export MAX_BATCHES="${MAX_BATCHES:-}"

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

LOGFILE="logs/visualize/visualize_probe_features_${PBS_JOBID}.log"

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

if [ ! -d "${PROBE_DIR}" ]; then
    echo "Error: PROBE_DIR not found: ${PROBE_DIR}" >&2
    exit 1
fi

if [ "${DATA_MODE}" = "toilet" ] && [ ! -d "${IMAGE_FOLDER}" ]; then
    echo "Error: IMAGE_FOLDER not found: ${IMAGE_FOLDER}" >&2
    exit 1
fi

if [ "${DATA_MODE}" = "cc3m" ] && [ -z "${HF_DATASET}" ] && [ -z "${DATA_DIR}" ]; then
    echo "Error: cc3m mode requires HF_DATASET+LOCAL_VAL_PATH or DATA_DIR." >&2
    exit 1
fi

if [ "${DATA_MODE}" = "coco" ] && [ -z "${HF_DATASET}" ]; then
    echo "Error: coco mode requires HF_DATASET+LOCAL_VAL_PATH." >&2
    exit 1
fi

if [ "${DATA_MODE}" = "folder" ] && [ -z "${DATA_DIR}" ]; then
    echo "Error: folder mode requires DATA_DIR." >&2
    exit 1
fi

# =============================================================================
# Setup
# =============================================================================
echo "========================================" | tee -a "${LOGFILE}"
echo "Job: Visualize Probe Features (Multi-Layer)" | tee -a "${LOGFILE}"
echo "Job ID:        ${PBS_JOBID}"              | tee -a "${LOGFILE}"
echo "Node:          $(hostname)"               | tee -a "${LOGFILE}"
echo "Started:       $(date)"                   | tee -a "${LOGFILE}"
echo "data_mode:     ${DATA_MODE}"              | tee -a "${LOGFILE}"
echo "sae_ckpt:      ${SAE_CKPT}"              | tee -a "${LOGFILE}"
echo "probe_dir:     ${PROBE_DIR}"             | tee -a "${LOGFILE}"
echo "layers:        ${LAYERS}"                | tee -a "${LOGFILE}"
echo "top_probe_k:   ${TOP_PROBE_K}"           | tee -a "${LOGFILE}"
echo "model_name:    ${MODEL_NAME}"             | tee -a "${LOGFILE}"
echo "dtype:         ${DTYPE}"                  | tee -a "${LOGFILE}"
echo "batch_size:    ${BATCH_SIZE}"             | tee -a "${LOGFILE}"
echo "output_dir:    ${OUTPUT_DIR}"             | tee -a "${LOGFILE}"
if [ "${DATA_MODE}" = "toilet" ]; then
    echo "image_folder:  ${IMAGE_FOLDER}"           | tee -a "${LOGFILE}"
    echo "toilet_mode:   ${OBJECT_MODE}"            | tee -a "${LOGFILE}"
    echo "caption_mode:  ${CAPTION_MODE}"           | tee -a "${LOGFILE}"
    echo "num_negatives: ${NUM_NEGATIVES}"          | tee -a "${LOGFILE}"
elif [ "${DATA_MODE}" = "cc3m" ]; then
    echo "hf_dataset:    ${HF_DATASET}"             | tee -a "${LOGFILE}"
    echo "local_val:     ${LOCAL_VAL_PATH}"         | tee -a "${LOGFILE}"
    echo "data_dir:      ${DATA_DIR}"               | tee -a "${LOGFILE}"
    echo "split:         ${SPLIT}"                  | tee -a "${LOGFILE}"
elif [ "${DATA_MODE}" = "coco" ]; then
    echo "hf_dataset:    ${HF_DATASET}"             | tee -a "${LOGFILE}"
    echo "local_val:     ${LOCAL_VAL_PATH}"         | tee -a "${LOGFILE}"
    echo "split:         ${SPLIT}"                  | tee -a "${LOGFILE}"
elif [ "${DATA_MODE}" = "folder" ]; then
    echo "data_dir:      ${DATA_DIR}"               | tee -a "${LOGFILE}"
fi
echo "========================================" | tee -a "${LOGFILE}"

# GPU check
echo "" | tee -a "${LOGFILE}"
echo "GPU Info:" | tee -a "${LOGFILE}"
nvidia-smi --query-gpu=index,name,memory.total --format=csv | tee -a "${LOGFILE}"

# =============================================================================
# Build argument list
# =============================================================================
ARGS=(
    --data_mode       "${DATA_MODE}"
    --sae_ckpt        "${SAE_CKPT}"
    --model_name      "${MODEL_NAME}"
    --dtype           "${DTYPE}"
    --probe_dir       "${PROBE_DIR}"
    --probe_input_dim "${PROBE_INPUT_DIM}"
    --layers          ${LAYERS}
    --top_probe_k     "${TOP_PROBE_K}"
    --num_workers     "${NUM_WORKERS}"
    --batch_size      "${BATCH_SIZE}"
    --sae_batch       "${SAE_BATCH}"
    --threshold       "${THRESHOLD}"
    --output_dir      "${OUTPUT_DIR}"
    --top_images      "${TOP_IMAGES}"
    --top_texts       "${TOP_TEXTS}"
    --buffer          "${BUFFER}"
)

# Data-mode-specific args
if [ "${DATA_MODE}" = "toilet" ]; then
    ARGS+=(
        --image_folder   "${IMAGE_FOLDER}"
        --object_mode    "${OBJECT_MODE}"
        --num_negatives  "${NUM_NEGATIVES}"
        --caption_mode   "${CAPTION_MODE}"
        --max_new_tokens "${MAX_NEW_TOKENS}"
    )
elif [ "${DATA_MODE}" = "cc3m" ]; then
    # cc3m supports either HF dataset + local path, or a plain image folder
    if [ -n "${HF_DATASET}" ]; then
        ARGS+=(--hf_dataset "${HF_DATASET}" --local_val_path "${LOCAL_VAL_PATH}" --split "${SPLIT}")
    else
        ARGS+=(--data_dir "${DATA_DIR}")
    fi
elif [ "${DATA_MODE}" = "coco" ]; then
    ARGS+=(--hf_dataset "${HF_DATASET}" --local_val_path "${LOCAL_VAL_PATH}" --split "${SPLIT}")
elif [ "${DATA_MODE}" = "folder" ]; then
    ARGS+=(--data_dir "${DATA_DIR}")
fi

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

# =============================================================================
# Run
# =============================================================================
echo "" | tee -a "${LOGFILE}"

if [ "${NUM_GPUS}" -gt 1 ]; then
    echo "Launching with torchrun on ${NUM_GPUS} GPUs..." | tee -a "${LOGFILE}"
    torchrun --nproc_per_node="${NUM_GPUS}" -m training.visualize_probe_features "${ARGS[@]}" \
        2>&1 | tee -a "${LOGFILE}"
else
    echo "Launching single-GPU mode..." | tee -a "${LOGFILE}"
    python -m training.visualize_probe_features "${ARGS[@]}" \
        2>&1 | tee -a "${LOGFILE}"
fi

EXIT_CODE=${PIPESTATUS[0]}

echo "" | tee -a "${LOGFILE}"
echo "========================================" | tee -a "${LOGFILE}"
if [ $EXIT_CODE -eq 0 ]; then
    echo "Status: SUCCESS" | tee -a "${LOGFILE}"
else
    echo "Status: FAILED (exit code $EXIT_CODE)" | tee -a "${LOGFILE}"
fi
echo "Completed: $(date)" | tee -a "${LOGFILE}"
echo "========================================" | tee -a "${LOGFILE}"

exit $EXIT_CODE