#!/bin/bash
# Per-layer bathroom↔toilet probe-readout correlation, PER MODEL VARIANT
# (base/lora/efuf/nullu), for four populations (bathroom_only, toilet_only,
# bathroom_toilet, random_cc3m/others). 4-variant probes; base caption reused
# across variants (one forward per variant).
#
# Tunables (override via env):
#   DEVICE_ID, DTYPE, NUM_IMAGES, SEED, HOOK_TYPE, MAX_NEW_TOKENS
#   METHODS       space-separated subset of: base lora efuf nullu (default all)
#   SCORE         logit (default) | sigma
#   BATH_PROBE / TOILET_PROBE   default = 4-variant probes
#   LORA_PATH / EFUF_PATH / NULLU_PATH / NULLU_LOWEST / NULLU_HIGHEST   edit ckpts
#   OUT, OUT_JSON output paths
#
# Usage:
#   bash mechanistic_interp/scripts/correlation_baselines.sh
#   NUM_IMAGES=100 DEVICE_ID=5 bash mechanistic_interp/scripts/correlation_baselines.sh

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

MODEL_NAME="llava-hf/llava-1.5-7b-hf"
DEVICE_ID="${DEVICE_ID:-0}"
DTYPE="${DTYPE:-bfloat16}"
NUM_IMAGES="${NUM_IMAGES:-100}"
SEED="${SEED:-0}"
METHODS="${METHODS:-base lora efuf nullu}"
SCORE="${SCORE:-logit}"
HOOK_TYPE="${HOOK_TYPE:-post}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-64}"

BATH_PROBE="${BATH_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_bathroom/post/seqprobe.pth}"
TOILET_PROBE="${TOILET_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_toilet/post/seqprobe.pth}"
LORA_PATH="${LORA_PATH:-/data/caotue/multilayer-sae/adv_gen_outputs/run_bathroom_toilet_v2/lora_adapter}"
EFUF_PATH="${EFUF_PATH:-/data/caotue/multilayer-sae/EFUF/efuf/checkpoints/llava_vicuna_7b/bathroom_toilet_paper_10ep/epoch_002.pth}"
NULLU_PATH="${NULLU_PATH:-/data/caotue/nullu/edited_models/LLaVA-7B-top4-0-32-bathroom_toilet}"
NULLU_LOWEST="${NULLU_LOWEST:-8}"
NULLU_HIGHEST="${NULLU_HIGHEST:-32}"

OUT="${OUT:-${REPO_ROOT}/mechanistic_interp/graph/correlation_baselines.png}"
OUT_JSON="${OUT_JSON:-${REPO_ROOT}/mechanistic_interp/graph/correlation_baselines.json}"

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

python -m mechanistic_interp.correlation_baselines \
    --model_name     "${MODEL_NAME}" \
    --device_id      "${DEVICE_ID}" \
    --dtype          "${DTYPE}" \
    --methods        ${METHODS} \
    --lora_path      "${LORA_PATH}" \
    --efuf_path      "${EFUF_PATH}" \
    --nullu_path     "${NULLU_PATH}" \
    --nullu_lowest   "${NULLU_LOWEST}" \
    --nullu_highest  "${NULLU_HIGHEST}" \
    --bath_probe     "${BATH_PROBE}" \
    --toilet_probe   "${TOILET_PROBE}" \
    --num_images     "${NUM_IMAGES}" \
    --seed           "${SEED}" \
    --score          "${SCORE}" \
    --hook_type      "${HOOK_TYPE}" \
    --max_new_tokens "${MAX_NEW_TOKENS}" \
    --out            "${OUT}" \
    --out_json       "${OUT_JSON}" \
    "$@"