#!/bin/bash
# =============================================================================
# integrated_gradient_reverse.sh — toilet→bathroom causal influence map via
# INTEGRATED GRADIENT of the toilet probe (REVERSED direction).
#
# Steer = toilet probe, readout = bathroom probe. Images: toilet-only (toilet=1,
# bathroom=0). Steered direction at layer l = integrated gradient of score_toilet_l
# along base→image, base = mean caption-token residual of NEGATIVE-label validation
# images. Step = ALPHA*‖h_l‖*ĝ plus a 'full' panel h'=h+(x−b). Δ bathroom score read
# at every layer l'>=l. Averaged → heatmap.
#
# Tunables (override inline, e.g. DEVICE_ID=2 IG_STEPS=64 bash <script>):
#   DEVICE_ID       GPU id                                  (default 0)
#   DTYPE           bfloat16 | float16 | float32            (default bfloat16)
#   BATH_PROBE      steered-direction checkpoint = TOILET probe in reverse
#   TOILET_PROBE    measured-readout  checkpoint = BATHROOM probe in reverse
#   IMAGE_FOLDER    folder of images
#   SAMPLES_JSON    samples.json path
#   BASE_PROMPT     prompt key in samples.json
#   STEER_NAME      display name of steered concept (plots)   (default toilet)
#   READOUT_NAME    display name of readout concept (plots)   (default bathroom)
#   NEG_JSONL       {"train":[ids],"validation":[ids]} negative stems
#   BASELINE_SPLIT  which NEG_JSONL split for the baseline    (default validation)
#   BASELINE_NUM    cap on negatives for baseline mean        (default 1000; 0=all)
#   IG_STEPS        Riemann steps for the IG path base→image  (default 32)
#   QUESTION        prompt question
#   NUM_IMAGES      cap on images; 0 = ALL toilet-only        (default 0)
#   ALPHAS          steps as FRACTION of ‖h_l‖                (default "0.05 0.1 0.2 0.4 0.8")
#   HOOK_TYPE       pre | mid | post                          (default post)
#   MAX_NEW_TOKENS  caption generation length                 (default 64)
#   MAX_SEQ_TOKENS  caption tokens kept for the probe         (default 64)
#   FORCED_TEXT     forced ASSISTANT answer  (default "In this image, there is a")
#   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}"
# REVERSED: BATH_PROBE = toilet probe (steer), TOILET_PROBE = bathroom probe (readout); 4variant
BATH_PROBE="${BATH_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_toilet/post/seqprobe.pth}"
TOILET_PROBE="${TOILET_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_bathroom/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.}"
STEER_NAME="${STEER_NAME:-toilet}"
READOUT_NAME="${READOUT_NAME:-bathroom}"
# Integrated-gradient baseline (mean NEGATIVE-label validation residual).
NEG_JSONL="${NEG_JSONL:-mechanistic_interp/neg_cc3m_5k.json}"
BASELINE_SPLIT="${BASELINE_SPLIT:-validation}"
BASELINE_NUM="${BASELINE_NUM:-1000}"
IG_STEPS="${IG_STEPS:-32}"
BASE_MENTIONS="${BASE_MENTIONS:-any}"
QUESTION="${QUESTION:-Describe this image.}"
FORCED_TEXT="${FORCED_TEXT:-In this image, there is a}"
NUM_IMAGES="${NUM_IMAGES:-0}"
ALPHAS="${ALPHAS:-0.05 0.1 0.2 0.4 0.8}"
HOOK_TYPE="${HOOK_TYPE:-post}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-64}"
MAX_SEQ_TOKENS="${MAX_SEQ_TOKENS:-64}"
OUT="${OUT:-mechanistic_interp/graph/integrated_gradient_toilet2bath.png}"
OUT_JSON="${OUT_JSON:-mechanistic_interp/graph/integrated_gradient_toilet2bath.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}"
    --bath_probe     "${BATH_PROBE}"
    --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}"
    --neg_jsonl      "${NEG_JSONL}"
    --baseline_split "${BASELINE_SPLIT}"
    --baseline_num   "${BASELINE_NUM}"
    --ig_steps       "${IG_STEPS}"
    --alphas         ${ALPHAS}
    --hook_type      "${HOOK_TYPE}"
    --max_new_tokens "${MAX_NEW_TOKENS}"
    --max_seq_tokens "${MAX_SEQ_TOKENS}"
    --steer_name     "${STEER_NAME}"
    --readout_name   "${READOUT_NAME}"
    --out            "${OUT}"
    --out_json       "${OUT_JSON}"
)
if [ -n "${FORCED_TEXT}" ]; then
    ARGS+=(--forced_text "${FORCED_TEXT}")
fi

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