#!/bin/bash
# =============================================================================
# random_gradient_ascent.sh — RANDOM-direction control for gradient_ascent.
#
# For each BATHROOM-ONLY image: steer the residual at layer l along a RANDOM unit
# direction (not the bathroom-probe gradient), step alpha*‖h_l‖*ĝ (or alpha*ĝ with
# HNORM=0), and measure Δ toilet score at every layer l'>=l. Averaged → heatmap.
# Null model for the bathroom→toilet effect: shows what any perturbation of
# comparable magnitude does.
#
# Tunables (override inline, e.g. NUM_IMAGES=100 HNORM=0 bash <script>):
#   DEVICE_ID       GPU id                                  (default 0)
#   DTYPE           bfloat16 | float16 | float32            (default bfloat16)
#   VARIANT         base | lora | nullu | efuf              (default base)
#   TOILET_PROBE    toilet SequenceLayerProbes (readout)
#   IMAGE_FOLDER    folder of images
#   SAMPLES_JSON    samples.json with base_mentions_object flags
#   BASE_MENTIONS   false | true | any                      (default false)
#   NUM_IMAGES      cap on images (0=ALL)                   (default 100)
#   SEED            RNG seed (image subsample + random dirs) (default 0)
#   ALPHAS          steps as FRACTION of ‖h_l‖ (or absolute if HNORM=0) (default "0.05 0.1 0.2 0.4 0.8")
#   HNORM           1 = scale step by ‖h_l‖ (default); 0 = --no-hnorm (absolute step)
#   HOOK_TYPE       pre | mid | post                        (default post)
#   PLOT_MODE       normal | meannorm                       (default meannorm)
#   FORCED_TEXT     forced ASSISTANT answer (empty = generated caption)
#   OUT / OUT_JSON  output heatmap PNG / matrix JSON
# =============================================================================

export HF_HOME="/data/caotue/hf_cache"
export HF_DATASETS_CACHE="/data/caotue/hf_cache/datasets"
export TORCH_HOME="/data/caotue/torch_cache"
export TMPDIR="/data/caotue/tmp"

DEVICE_ID="${DEVICE_ID:-0}"
DTYPE="${DTYPE:-bfloat16}"
VARIANT="${VARIANT:-base}"
TOILET_PROBE="${TOILET_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_toilet/post/seqprobe.pth}"
IMAGE_FOLDER="${IMAGE_FOLDER:-/data/caotue/CC3M-Dataset/cc3m_images}"
SAMPLES_JSON="${SAMPLES_JSON:-mechanistic_interp/toilet_bathroom/samples.json}"
BASE_PROMPT="${BASE_PROMPT:-Describe this image.}"
BASE_MENTIONS="${BASE_MENTIONS:-false}"
QUESTION="${QUESTION:-Describe this image.}"
FORCED_TEXT="${FORCED_TEXT:-This image features a bathroom with a}"
NUM_IMAGES="${NUM_IMAGES:-100}"
SEED="${SEED:-0}"
ALPHAS="${ALPHAS:-0.05 0.1 0.2 0.4 0.8}"
HNORM="${HNORM:-1}"
HOOK_TYPE="${HOOK_TYPE:-post}"
PLOT_MODE="${PLOT_MODE:-meannorm}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-64}"
MAX_SEQ_TOKENS="${MAX_SEQ_TOKENS:-64}"
OUT="${OUT:-mechanistic_interp/graph/random_gradient_ascent_bath2toilet.png}"
OUT_JSON="${OUT_JSON:-mechanistic_interp/graph/random_gradient_ascent_bath2toilet.json}"

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

ARGS=(
    --device_id      "${DEVICE_ID}"
    --dtype          "${DTYPE}"
    --variant        "${VARIANT}"
    --toilet_probe   "${TOILET_PROBE}"
    --image_folder   "${IMAGE_FOLDER}"
    --samples_json   "${SAMPLES_JSON}"
    --base_prompt    "${BASE_PROMPT}"
    --base_mentions  "${BASE_MENTIONS}"
    --question       "${QUESTION}"
    --num_images     "${NUM_IMAGES}"
    --seed           "${SEED}"
    --alphas         ${ALPHAS}
    --hook_type      "${HOOK_TYPE}"
    --plot_mode      "${PLOT_MODE}"
    --max_new_tokens "${MAX_NEW_TOKENS}"
    --max_seq_tokens "${MAX_SEQ_TOKENS}"
    --out            "${OUT}"
    --out_json       "${OUT_JSON}"
)
# ‖h‖ scaling toggle: HNORM=0 → --no-hnorm (absolute step).
if [ "${HNORM}" = "0" ] || [ "${HNORM}" = "false" ]; then
    ARGS+=(--no-hnorm)
fi
# Only pass --forced_text when non-empty (empty = use generated captions).
if [ -n "${FORCED_TEXT}" ]; then
    ARGS+=(--forced_text "${FORCED_TEXT}")
fi

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