#!/bin/bash
# =============================================================================
# probe_output_trace.sh — per-layer probe P(toilet) under fixed-prefix
# forward passes, one curve per method (base / ours / Nullu / EFUF).
# =============================================================================
set -euo pipefail
cd "$(dirname "$0")/../.."
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"

PROBE_CKPT="${PROBE_CKPT:-/mnt/raid10/erwin/testing/multilayer-sae/full_runs/p2_20260511_024729/probes_gen_bathroom_toilet.pt}"
SAE_CKPT="${SAE_CKPT:-/mnt/raid10/erwin/testing/multilayer-sae/last.ckpt}"
SAMPLES_JSON="${SAMPLES_JSON:-mechanistic_interp/toilet-bathroom/lora_adapter/samples.json}"
PROMPT="${PROMPT:-Describe this image.}"
HF_DATASET="${HF_DATASET:-pbcong/bathroom-toilet}"
HF_SPLIT="${HF_SPLIT:-validation}"
ID_COL="${ID_COL:-image_id}"

ADAPTER_PATH="${ADAPTER_PATH:-/mnt/raid10/erwin/testing/multilayer-sae/adv_gen_outputs/p2_20260511_024729/run_p2_20260511_024729/lora_adapter/adapter_model.safetensors}"
ADAPTER_CFG="${ADAPTER_CFG:-/mnt/raid10/erwin/testing/multilayer-sae/adv_gen_outputs/p2_20260511_024729/run_p2_20260511_024729/lora_adapter/adapter_config.json}"

MODEL_NAME="${MODEL_NAME:-llava-hf/llava-1.5-7b-hf}"
DEVICE="${DEVICE:-cuda:0}"
DTYPE="${DTYPE:-bfloat16}"

N_SAMPLES="${N_SAMPLES:-}"             # empty → all category-matching samples
HOOK_TYPE="${HOOK_TYPE:-post}"
POOL="${POOL:-max}"                    # match the pool used at probe training
CATEGORY="${CATEGORY:-bathroom_only}"

FIXED_ASSISTANT_PREFIX="${FIXED_ASSISTANT_PREFIX:-In this bathroom there is a shower, a sink and}"
if [[ -z "${FIXED_ASSISTANT_PREFIX}" ]]; then
    echo "FIXED_ASSISTANT_PREFIX must be non-empty for probe_output_trace.sh" >&2
    exit 1
fi
PREFIX_SLUG="$(echo "${FIXED_ASSISTANT_PREFIX}" \
               | tr '[:upper:]' '[:lower:]' \
               | sed -e 's/[^a-z0-9]\+/_/g' -e 's/^_//' -e 's/_$//' \
               | cut -c1-40)"
PREFIX_TAG="fixed_${PREFIX_SLUG}"
PROBE_TAG="${PROBE_TAG:-$(basename "${PROBE_CKPT}" .pt)}"

OUT_DIR="${OUT_DIR:-mechanistic_interp/probe_output_trace/${CATEGORY}/${POOL}/${PROBE_TAG}/${PREFIX_TAG}/traces}"
GRAPH_DIR="${GRAPH_DIR:-mechanistic_interp/probe_output_trace/${CATEGORY}/${POOL}/${PROBE_TAG}/${PREFIX_TAG}/graphs}"

NULLU_MODEL_PATH="${NULLU_MODEL_PATH:-/mnt/raid10/erwin/testing/multilayer-sae/Nullu/output/edited_model/LLaVA-7B-top4-0-32-bathroom_toilet}"
NULLU_LOWEST_LAYER="${NULLU_LOWEST_LAYER:-8}"
NULLU_HIGHEST_LAYER="${NULLU_HIGHEST_LAYER:-32}"
EFUF_CKPT="${EFUF_CKPT:-/mnt/raid10/erwin/testing/multilayer-sae/EFUF/efuf/checkpoints/llava_vicuna_7b/bathroom_toilet_20ep_fixed/epoch_019.pth}"

ARGS=(
    -m mechanistic_interp.probe_output_trace
    --probe_ckpt          "${PROBE_CKPT}"
    --sae_ckpt            "${SAE_CKPT}"
    --samples_json        "${SAMPLES_JSON}"
    --prompt              "${PROMPT}"
    --hf_dataset          "${HF_DATASET}"
    --hf_split            "${HF_SPLIT}"
    --id_col              "${ID_COL}"
    --adapter_path        "${ADAPTER_PATH}"
    --adapter_cfg         "${ADAPTER_CFG}"
    --model_name          "${MODEL_NAME}"
    --device              "${DEVICE}"
    --dtype               "${DTYPE}"
    --hook_type           "${HOOK_TYPE}"
    --pool                "${POOL}"
    --category            "${CATEGORY}"
    --fixed_assistant_prefix "${FIXED_ASSISTANT_PREFIX}"
    --out_dir             "${OUT_DIR}"
    --graph_dir           "${GRAPH_DIR}"
)
if [[ -n "${N_SAMPLES}" ]]; then
    ARGS+=(--n_samples "${N_SAMPLES}")
fi
if [[ -n "${NULLU_MODEL_PATH}" ]]; then
    ARGS+=(
        --nullu_model_path    "${NULLU_MODEL_PATH}"
        --nullu_lowest_layer  "${NULLU_LOWEST_LAYER}"
        --nullu_highest_layer "${NULLU_HIGHEST_LAYER}"
    )
fi
if [[ -n "${EFUF_CKPT}" ]]; then
    ARGS+=(--efuf_ckpt "${EFUF_CKPT}")
fi

PY="${PY:-/mnt/raid10/erwin/testing/multilayer-sae/.venv/bin/python}"
echo "Running: $PY ${ARGS[*]}"
"$PY" "${ARGS[@]}"