#!/bin/bash

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

# =============================================================================
# PBS Job File: Compare SAE Activations — Original vs Modified LLaVA
#
# Runs training/compare_activations.py which:
#   1. Loads linear probe checkpoints → top-k features + sub-threshold features
#   2. Runs forward passes through both models on the same inputs
#   3. Computes per-feature activation stats for top-k features
#   4. Computes |diff| and JSD for sub-threshold features
#   5. Saves comparison.json
#
# FILTER MODES  (set FILTER_MODE below)
# ------------
#   bathroom — bathroom == 1  (regardless of toilet value)
#   toilet   — toilet   == 1  (regardless of bathroom value)
#   both     — toilet   == 1 OR bathroom == 1
#   all      — no filtering
#
# To submit:
#   qsub training/scripts/run_compare_activations.pbs
#
# Override any variable:
#   FILTER_MODE=toilet MODIFIED_MODEL=/path/to/model \
#   qsub training/scripts/run_compare_activations.pbs
#
# To add a third (Nullu-edited) model to the comparison, set NULLU_MODEL_PATH
# to the directory produced by Nullu's experiments/llava_edit.pbs, and optionally
# override NULLU_LOWEST_LAYER / NULLU_HIGHEST_LAYER:
#   NULLU_MODEL_PATH=/path/to/Nullu/output/edited_model/LLaVA-7B-top4-0-32-test \
#   NULLU_LOWEST_LAYER=16 NULLU_HIGHEST_LAYER=32 \
#   qsub training/scripts/run_compare_activations.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/compare

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

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

# ── Models ────────────────────────────────────────────────────────────────────
export ORIGINAL_MODEL="${ORIGINAL_MODEL:-llava-hf/llava-1.5-7b-hf}"
export MODIFIED_MODEL="${MODIFIED_MODEL:-/home/users/ntu/cong045/scratch/testing/multilayer-sae/adv_outputs/bathroom_toilet/run_20260405_152304/lora_adapter}"
export SAE_CKPT="${SAE_CKPT:-training/multilayer_sae_ckpt/last.ckpt}"
export DTYPE="${DTYPE:-float16}"

# ── Nullu model (optional; leave NULLU_MODEL_PATH empty to skip) ──────────────
# Point NULLU_MODEL_PATH at Nullu's edited checkpoint directory, e.g.
# /path/to/Nullu/output/edited_model/LLaVA-7B-top4-0-32-test
export NULLU_MODEL_PATH="${NULLU_MODEL_PATH:-}"
export NULLU_LOWEST_LAYER="${NULLU_LOWEST_LAYER:-16}"
export NULLU_HIGHEST_LAYER="${NULLU_HIGHEST_LAYER:-32}"

# ── 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:-1000}"   # features tracked & saved to JSON
export TOP_PRINT_K="${TOP_PRINT_K:-20}"    # features printed per-row in log

# ── Data ──────────────────────────────────────────────────────────────────────
export DATA_MODE="${DATA_MODE:-single}"          # single | toilet | cc3m | coco
export CAPTION_MODE="${CAPTION_MODE:-generated}" # generated | caption

# [toilet] HF dataset + local image folder + bathroom/toilet filter
export HF_DATASET="${HF_DATASET:-pbcong/bathroom-toilet}"
export IMAGE_FOLDER="${IMAGE_FOLDER:-/home/users/ntu/cong045/scratch/testing/hallucination/CC3M-Dataset/cc3m_images/train}"
export FILTER_MODE="${FILTER_MODE:-bathroom}"    # bathroom | toilet | both | all

# [single] Path to one image file
export IMAGE_PATH="${IMAGE_PATH:-/home/users/ntu/cong045/scratch/testing/hallucination/CC3M-Dataset/cc3m_images/train/000021450.jpg}"

# [cc3m / coco] Local image root + dataset split
export LOCAL_VAL_PATH="${LOCAL_VAL_PATH:-}"
export SPLIT="${SPLIT:-train}"

# ── Distribution measure ───────────────────────────────────────────────────────
export MEASURE_DISTRIBUTION="${MEASURE_DISTRIBUTION:-kl_divergence}"  # jensen_shannon_divergence | kl_divergence

# ── Diff norm ─────────────────────────────────────────────────────────────────
export DIFF_NORM="${DIFF_NORM:-l1}"  # l1 | l2

# ── Processing ────────────────────────────────────────────────────────────────
export BATCH_SIZE="${BATCH_SIZE:-64}"
export SAE_BATCH="${SAE_BATCH:-4096}"
export NUM_WORKERS="${NUM_WORKERS:-16}"
export MAX_BATCHES="${MAX_BATCHES:-}"
export DIST_TIMEOUT_HOURS="${DIST_TIMEOUT_HOURS:-8}"  # NCCL timeout; increase if batches take >1h

# ── Output ────────────────────────────────────────────────────────────────────
export OUTPUT_DIR="${OUTPUT_DIR:-training/compare_activations}"

LOGFILE="logs/compare/compare_activations_${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}" = "single" ]; then
    if [ -z "${IMAGE_PATH}" ]; then
        echo "Error: DATA_MODE=single requires IMAGE_PATH" >&2
        exit 1
    fi
    if [ ! -f "${IMAGE_PATH}" ]; then
        echo "Error: IMAGE_PATH not found: ${IMAGE_PATH}" >&2
        exit 1
    fi
elif [ "${DATA_MODE}" = "toilet" ]; then
    if [ ! -d "${IMAGE_FOLDER}" ]; then
        echo "Error: IMAGE_FOLDER not found: ${IMAGE_FOLDER}" >&2
        exit 1
    fi
elif [ "${DATA_MODE}" = "cc3m" ] || [ "${DATA_MODE}" = "coco" ]; then
    if [ -z "${LOCAL_VAL_PATH}" ]; then
        echo "Error: DATA_MODE=${DATA_MODE} requires LOCAL_VAL_PATH" >&2
        exit 1
    fi
else
    echo "Error: Unknown DATA_MODE=${DATA_MODE} (single | toilet | cc3m | coco)" >&2
    exit 1
fi

# =============================================================================
# Setup
# =============================================================================
echo "========================================" | tee -a "${LOGFILE}"
echo "Job: Compare Activations (Original vs Modified)" | tee -a "${LOGFILE}"
echo "Job ID:           ${PBS_JOBID}"              | tee -a "${LOGFILE}"
echo "Node:             $(hostname)"               | tee -a "${LOGFILE}"
echo "Started:          $(date)"                   | tee -a "${LOGFILE}"
echo "original_model:   ${ORIGINAL_MODEL}"         | tee -a "${LOGFILE}"
echo "modified_model:   ${MODIFIED_MODEL}"         | tee -a "${LOGFILE}"
if [ -n "${NULLU_MODEL_PATH}" ]; then
    echo "nullu_model:      ${NULLU_MODEL_PATH}"    | tee -a "${LOGFILE}"
    echo "nullu_layers:     ${NULLU_LOWEST_LAYER}-${NULLU_HIGHEST_LAYER}" \
                                                    | tee -a "${LOGFILE}"
fi
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 "top_print_k:       ${TOP_PRINT_K}"                 | tee -a "${LOGFILE}"
echo "data_mode:         ${DATA_MODE}"                   | tee -a "${LOGFILE}"
echo "caption_mode:      ${CAPTION_MODE}"                | tee -a "${LOGFILE}"
echo "measure_dist:      ${MEASURE_DISTRIBUTION}"        | tee -a "${LOGFILE}"
echo "diff_norm:         ${DIFF_NORM}"                  | tee -a "${LOGFILE}"
if [ "${DATA_MODE}" = "single" ]; then
    echo "image_path:    ${IMAGE_PATH}"                  | tee -a "${LOGFILE}"
elif [ "${DATA_MODE}" = "toilet" ]; then
    echo "hf_dataset:    ${HF_DATASET}"                  | tee -a "${LOGFILE}"
    echo "image_folder:  ${IMAGE_FOLDER}"                | tee -a "${LOGFILE}"
    echo "filter_mode:   ${FILTER_MODE}"                 | tee -a "${LOGFILE}"
elif [ "${DATA_MODE}" = "cc3m" ] || [ "${DATA_MODE}" = "coco" ]; then
    echo "hf_dataset:    ${HF_DATASET}"                  | tee -a "${LOGFILE}"
    echo "local_val_path:${LOCAL_VAL_PATH}"              | tee -a "${LOGFILE}"
    echo "split:         ${SPLIT}"                       | tee -a "${LOGFILE}"
fi
echo "batch_size:        ${BATCH_SIZE}"                  | 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=(
    --original_model       "${ORIGINAL_MODEL}"
    --modified_model       "${MODIFIED_MODEL}"
    --sae_ckpt             "${SAE_CKPT}"
    --dtype                "${DTYPE}"
    --probe_dir            "${PROBE_DIR}"
    --probe_input_dim      "${PROBE_INPUT_DIM}"
    --layers               ${LAYERS}
    --top_probe_k          "${TOP_PROBE_K}"
    --top_print_k          "${TOP_PRINT_K}"
    --data_mode            "${DATA_MODE}"
    --caption_mode         "${CAPTION_MODE}"
    --measure_distribution "${MEASURE_DISTRIBUTION}"
    --diff_norm            "${DIFF_NORM}"
    --batch_size           "${BATCH_SIZE}"
    --sae_batch            "${SAE_BATCH}"
    --num_workers          "${NUM_WORKERS}"
    --output_dir           "${OUTPUT_DIR}"
    --dist_timeout_hours   "${DIST_TIMEOUT_HOURS}"
)

# Data-mode-specific arguments
if [ "${DATA_MODE}" = "single" ]; then
    ARGS+=(--image_path "${IMAGE_PATH}")
elif [ "${DATA_MODE}" = "toilet" ]; then
    ARGS+=(--hf_dataset   "${HF_DATASET}")
    ARGS+=(--image_folder "${IMAGE_FOLDER}")
    ARGS+=(--filter_mode  "${FILTER_MODE}")
elif [ "${DATA_MODE}" = "cc3m" ] || [ "${DATA_MODE}" = "coco" ]; then
    ARGS+=(--hf_dataset     "${HF_DATASET}")
    ARGS+=(--local_val_path "${LOCAL_VAL_PATH}")
    ARGS+=(--split          "${SPLIT}")
fi

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

if [ -n "${NULLU_MODEL_PATH}" ]; then
    ARGS+=(--nullu_model_path    "${NULLU_MODEL_PATH}")
    ARGS+=(--nullu_lowest_layer  "${NULLU_LOWEST_LAYER}")
    ARGS+=(--nullu_highest_layer "${NULLU_HIGHEST_LAYER}")
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.compare_activations "${ARGS[@]}" \
        2>&1 | tee -a "${LOGFILE}"
else
    echo "Launching single-GPU mode..." | tee -a "${LOGFILE}"
    python -m training.compare_activations "${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