#!/bin/bash
# SAE feature verification / patching driver (mechanistic_interp/sae_patching.py).
#
# DEFAULT MODE = validate: the only mode that actually CONFIRMS a feature set.
# For each (src=concept image, tgt=clean image) pair and each layer it runs:
#   insertion (SUFFICIENCY): add the src selected-feature contribution into the
#                            clean image -> the concept token should go UP.
#   ablation  (NECESSITY):   remove the selected features from the concept image
#                            -> the concept token should go DOWN.
# Each is compared against N_RANDOM random feature sets (matched count) and
# averaged over all image pairs. Confirmed when, at the responsible layers:
#   insertion: selected >> random > 0   AND   ablation: selected << random < 0.
# The SAE is the trained checkpoint loaded as TopKSAE (per-token top-k=8).
#
# SOURCE images default to SAMPLES_N images from samples.json where SAMPLES_COL==1
# (concept present, e.g. bathroom), each paired with the single clean TGT_IMAGE
# (cat). Set SAMPLES="" to instead supply explicit SRC_IMAGES/TGT_IMAGES lists.
#
# Tunables (override via env):
#   MODE           validate | feature_patch | causal | steering   (default validate)
#   FEATURES_JSON  per-layer feature set to verify (bathroom or toilet)
#   TOKEN          concept token to measure (bath; toilet 1st subword is to)
#   PREFIX         teacher-forced text after "ASSISTANT:" (next-token = TOKEN)
#   SAMPLES        samples.json to draw concept-present sources from (""=use lists)
#   SAMPLES_COL    samples.json column marking concept-present ==1  (default bathroom)
#   SAMPLES_N      number of source images to take                  (default 20)
#   IMAGE_ROOT     dir holding <image_id>.jpg for SAMPLES
#   TGT_IMAGE      single clean target image, reused per source     (default cat.png)
#   SRC_IMAGES/TGT_IMAGES  explicit space-separated lists (used only when SAMPLES="")
#   N_RANDOM       random-feature control draws per layer (default 20; >50 is waste)
#   SEED, SAE_CKPT, HOOK_TYPE, DEVICE, DTYPE, OUTPUT_PREFIX
#
# Usage:
#   bash mechanistic_interp/scripts/sae_patching.sh                              # bathroom: 20 src -> cat
#   SAMPLES_N=40 bash mechanistic_interp/scripts/sae_patching.sh                 # more source images
#   FEATURES_JSON=.../toilet_top20.json TOKEN=to SAMPLES_COL=toilet \
#       bash mechanistic_interp/scripts/sae_patching.sh                          # verify toilet feats
#   SAMPLES="" SRC_IMAGES="a.jpg b.jpg" TGT_IMAGES="c.png d.png" \
#       bash mechanistic_interp/scripts/sae_patching.sh                          # explicit pairs

set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"

MODEL_NAME="${MODEL_NAME:-llava-hf/llava-1.5-7b-hf}"
SAE_CKPT="${SAE_CKPT:-/data/caotue/SAE_checkpoints/16_d_model/last.ckpt}"
FEATURES_JSON="${FEATURES_JSON:-${REPO_ROOT}/mechanistic_interp/probes/selected_features_f1_post_toilet_top10.json}"

MODE="${MODE:-validate}"
TOKEN="${TOKEN:-toilet}"          # bathroom: bath ; toilet: use PREFIX '...a to' + TOKEN_ID=488 ('ile')
TOKEN_ID="${TOKEN_ID:-488}"     # optional exact vocab id (overrides TOKEN); e.g. 488 = 'ile' in 'toilet'
PREFIX="${PREFIX:-This image features a to}"
HOOK_TYPE="${HOOK_TYPE:-post}"
DEVICE="${DEVICE:-cuda:4}"
DTYPE="${DTYPE:-bfloat16}"
N_RANDOM="${N_RANDOM:-20}" # random features
SEED="${SEED:-0}"
PATCH_BATCH="${PATCH_BATCH:-16}"   # feature sets patched per forward (higher=faster, more VRAM)

# SOURCE = SAMPLES_N concept-present images from samples.json (SAMPLES_COL==1).
# TARGET = one clean image (cat), reused for every source.
SAMPLES="${SAMPLES:-${REPO_ROOT}/mechanistic_interp/toilet_bathroom/samples.json}"
SAMPLES_COL="${SAMPLES_COL:-toilet}" # bathroom, toilet
SAMPLES_N="${SAMPLES_N:-20}"

# Output name defaults to validate_<column> (e.g. validate_bathroom, validate_toilet)
# -> ${OUTPUT_PREFIX}_results.json and ${OUTPUT_PREFIX}_validate.png.
OUTPUT_PREFIX="${OUTPUT_PREFIX:-mechanistic_interp/graph/sae_patch/validate_${SAMPLES_COL}}"
IMAGE_ROOT="${IMAGE_ROOT:-/data/caotue/CC3M-Dataset/cc3m_images/train}"
TGT_IMAGE="${TGT_IMAGE:-${REPO_ROOT}/mechanistic_interp/toilet_bathroom/cat.png}"

# Explicit pairs, used only when SAMPLES="" (paired by index).
SRC_IMAGES="${SRC_IMAGES:-${IMAGE_ROOT}/001607409.jpg}"
TGT_IMAGES="${TGT_IMAGES:-${TGT_IMAGE}}"

# ------------------------------------------------------------------ #
mkdir -p "$(dirname "${OUTPUT_PREFIX}")"

echo "========================================================"
echo "  SAE feature verification  (mode=${MODE})"
echo "========================================================"
echo "  Features   : ${FEATURES_JSON}"
echo "  Token      : 'ASSISTANT:${PREFIX}' -> measuring '${TOKEN}'"
if [ -n "${SAMPLES}" ]; then
echo "  Source     : ${SAMPLES_N} images from ${SAMPLES} (${SAMPLES_COL}==1)"
echo "  Target     : ${TGT_IMAGE} (clean, reused per source)"
else
echo "  Src images : ${SRC_IMAGES}"
echo "  Tgt images : ${TGT_IMAGES}"
fi
echo "  N random   : ${N_RANDOM} | seed ${SEED}"
echo "  SAE ckpt   : ${SAE_CKPT} (loaded as TopKSAE, top-k=8)"
echo "  Device     : ${DEVICE} / ${DTYPE} | hook ${HOOK_TYPE}"
echo "  Output     : ${OUTPUT_PREFIX}_{validate.png,results.json}"
echo "========================================================"

cd "${REPO_ROOT}"
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"

ARGS=(
    --model_name    "${MODEL_NAME}"
    --mode          "${MODE}"
    --tgt_image     "${TGT_IMAGE}"
    --prefix        "${PREFIX}"
    --token         "${TOKEN}"
    --hook_type     "${HOOK_TYPE}"
    --device        "${DEVICE}"
    --dtype         "${DTYPE}"
    --sae_ckpt      "${SAE_CKPT}"
    --features_json "${FEATURES_JSON}"
    --n_random      "${N_RANDOM}"
    --seed          "${SEED}"
    --patch_batch   "${PATCH_BATCH}"
    --output_prefix "${OUTPUT_PREFIX}"
)

[ -n "${TOKEN_ID}" ] && ARGS+=(--token_id "${TOKEN_ID}")

if [ -n "${SAMPLES}" ]; then
    # samples mode: build sources from samples.json; --src_image is a required
    # placeholder (unused in samples mode), so point it at the clean target.
    ARGS+=(
        --src_image   "${TGT_IMAGE}"
        --samples     "${SAMPLES}"
        --samples_col "${SAMPLES_COL}"
        --samples_n   "${SAMPLES_N}"
        --image_root  "${IMAGE_ROOT}"
    )
else
    ARGS+=(
        --src_image  "${SRC_IMAGES%% *}"
        --src_images ${SRC_IMAGES}
        --tgt_images ${TGT_IMAGES}
    )
fi

python -m mechanistic_interp.confirm_features_patching "${ARGS[@]}" "$@"