#!/bin/bash
# =============================================================================
# knowledge_suppression_trace.sh — matched-input internal-activation trace.
#
# For each bathroom-only image (object absent, scene present):
#   1) Base free-greedy-decodes K tokens — the matched input.
#   2) Each method (base / ours / Nullu) is teacher-forced on that K-token
#      continuation.
#   3) At every selected layer, residuals at the K generated text positions
#      are SAE-encoded and the pre-selected confident toilet features are
#      gathered. Per-layer scalar = max (or mean) over (K × top-k).
#   4) Per-image PNG + population summary across all images.
#
# Reads `FEATURES_JSON` for the per-layer confident toilet features.
# =============================================================================
set -euo pipefail
cd "$(dirname "$0")/../.."
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"

FEATURES_JSON="${FEATURES_JSON:-mechanistic_interp/probes/selected_features_probe_top10.json}"
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:-mechanistic_interp/toilet-bathroom/lora_adapter/adapter_model.safetensors}"
ADAPTER_CFG="${ADAPTER_CFG:-mechanistic_interp/toilet-bathroom/lora_adapter/adapter_config.json}"
SAE_CKPT="${SAE_CKPT:-training/multilayer_sae_ckpt/last.ckpt}"

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

N_SAMPLES="${N_SAMPLES:-}"             # empty → all bathroom_only samples
GEN_TOKENS="${GEN_TOKENS:-32}"
HOOK_TYPE="${HOOK_TYPE:-post}"
SAE_BATCH="${SAE_BATCH:-2048}"
CATEGORY="${CATEGORY:-bathroom_only}"  # D_{A,¬c} from the spec
AGG="${AGG:-max}"                       # max | mean

# Deterministic assistant-side prefix. Empty → base free-decodes K tokens.
# Non-empty → use this fixed string (same for every method); K = its token count.
FIXED_ASSISTANT_PREFIX="${FIXED_ASSISTANT_PREFIX:-}"
if [[ -n "${FIXED_ASSISTANT_PREFIX}" ]]; then
    # Sluggify the prefix (alnum+underscore, lowercase, capped at 40 chars) so
    # different prefixes don't overwrite each other.
    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}"
else
    PREFIX_TAG="freegen"
fi

PROBE_TAG="${PROBE_TAG:-$(basename "${FEATURES_JSON}" .json)}"
OUT_DIR="${OUT_DIR:-mechanistic_interp/knowledge_suppression/${CATEGORY}/${AGG}/${PROBE_TAG}/${PREFIX_TAG}/traces}"
GRAPH_DIR="${GRAPH_DIR:-mechanistic_interp/knowledge_suppression/${CATEGORY}/${AGG}/${PROBE_TAG}/${PREFIX_TAG}/graphs}"

# Nullu (optional). Empty NULLU_MODEL_PATH → Nullu curve omitted.
NULLU_MODEL_PATH="${NULLU_MODEL_PATH:-}"
NULLU_LOWEST_LAYER="${NULLU_LOWEST_LAYER:-16}"
NULLU_HIGHEST_LAYER="${NULLU_HIGHEST_LAYER:-32}"

# EFUF (optional). Empty EFUF_CKPT → EFUF curve omitted.
EFUF_CKPT="${EFUF_CKPT:-}"

ARGS=(
    -m mechanistic_interp.knowledge_suppression_trace
    --features_json     "${FEATURES_JSON}"
    --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}"
    --sae_ckpt          "${SAE_CKPT}"
    --model_name        "${MODEL_NAME}"
    --device            "${DEVICE}"
    --dtype             "${DTYPE}"
    --gen_tokens        "${GEN_TOKENS}"
    --hook_type         "${HOOK_TYPE}"
    --sae_batch         "${SAE_BATCH}"
    --category          "${CATEGORY}"
    --agg               "${AGG}"
    --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
if [[ -n "${FIXED_ASSISTANT_PREFIX}" ]]; then
    ARGS+=(--fixed_assistant_prefix "${FIXED_ASSISTANT_PREFIX}")
fi

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