#!/bin/bash

#PBS -N visualize_sae_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 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).
#
# To submit:
#   qsub training/scripts/run_visualize_multilayer_sae_features.pbs
#
# Override variables before submission:
#   DATA_MODE=cc3m HF_DATASET=pixparse/cc3m-wds \
#   LOCAL_VAL_PATH=/path/to/cc3m/train \
#   FEATURE_IDS="0 1 42" \
#   qsub training/scripts/run_visualize_multilayer_sae_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}"

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

# Mode 1:
# ── Data — toilet ─────────────────────────────────────────────────────────────
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 NUM_NEGATIVES="${NUM_NEGATIVES:-10000}"

# Mode 2 + 3:
# ── Data — CC3M / COCO / folder ───────────────
export HF_DATASET="${HF_DATASET:-}"              # e.g. pixparse/cc3m-wds | yerevann/coco-karpathy
export LOCAL_VAL_PATH="${LOCAL_VAL_PATH:-}"      # local image root for HF dataset
export DATA_DIR="${DATA_DIR:-}"                  # [folder mode] plain image directory
export SPLIT="${SPLIT:-train}"

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

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

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

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

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

# =============================================================================
# Validation
# =============================================================================
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 (e.g. FEATURE_IDS=\"0 1 42\")." >&2
    exit 1
fi

# =============================================================================
# Setup
# =============================================================================
echo "========================================" | tee -a "${LOGFILE}"
echo "Job: Visualize SAE Features"              | tee -a "${LOGFILE}"
echo "Job ID:        ${PBS_JOBID}"              | tee -a "${LOGFILE}"
echo "Node:          $(hostname)"               | tee -a "${LOGFILE}"
echo "Started:       $(date)"                   | tee -a "${LOGFILE}"
echo "sae_ckpt:      ${SAE_CKPT}"              | tee -a "${LOGFILE}"
echo "model_name:    ${MODEL_NAME}"             | tee -a "${LOGFILE}"
echo "data_mode:     ${DATA_MODE}"              | tee -a "${LOGFILE}"
echo "caption_mode:  ${CAPTION_MODE}"           | tee -a "${LOGFILE}"
echo "image_folder:  ${IMAGE_FOLDER}"           | tee -a "${LOGFILE}"
echo "toilet_mode:   ${OBJECT_MODE}"            | tee -a "${LOGFILE}"
echo "hf_dataset:    ${HF_DATASET}"             | tee -a "${LOGFILE}"
echo "hook_point:    ${HOOK_POINT}"             | tee -a "${LOGFILE}"
echo "feature_ids:   ${FEATURE_IDS}"            | tee -a "${LOGFILE}"
echo "batch_size:    ${BATCH_SIZE}"             | tee -a "${LOGFILE}"
echo "dtype:         ${DTYPE}"                  | tee -a "${LOGFILE}"
echo "output_dir:    ${OUTPUT_DIR}"             | tee -a "${LOGFILE}"
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=(
    --sae_ckpt        "${SAE_CKPT}"
    --model_name      "${MODEL_NAME}"
    --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}"
            --object_mode   "${OBJECT_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
# =============================================================================
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_multilayer_sae_features "${ARGS[@]}" \
        2>&1 | tee -a "${LOGFILE}"
else
    echo "Launching single-GPU mode..." | tee -a "${LOGFILE}"
    python -m training.visualize_multilayer_sae_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