#!/bin/bash
# =============================================================================
# per_layer_effect.sh — fixed-prefix probe per intervention layer.
#
# For each filter-passing sample, teacher-force `prompt + " " + fixed_assistant_prefix`
# (no decode) under base, then per-L_int ΔW, then ΔW@all-layers. At every
# capture layer L_cap: residual @ K prefix positions → FrozenSAEEncoder →
# pool over K → gather selected features + apply layer-L probe.
#
# Output per sample (under ${GRAPH_DIR}/${image_id}/):
#   * per_layer_effect_L*.png         — feature grids (L_int in groups of 4)
#   * per_layer_effect_probe_L*.png   — probe-output grids (same grouping)
#   * per_layer_effect_all.png        — base vs ΔW@all summary (features + probes)
# Plus ${OUT_DIR}/${image_id}.pt per sample with the raw tensors.
# =============================================================================
set -euo pipefail
cd "$(dirname "$0")/../.."
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"

# Object names for the dataset
OBJECT_NAME="${OBJECT_NAME:-toilet}"
OBJECT2_NAME="${OBJECT2_NAME:-bathroom}"

FEATURES_JSON="${FEATURES_JSON:-mechanistic_interp/probes/selected_features_probe_top20.json}"
SAMPLES_JSON="${SAMPLES_JSON:-mechanistic_interp/toilet-bathroom/lora_adapter/samples.json}"
PROMPT="${PROMPT:-Describe this image.}"
FIXED_ASSISTANT_PREFIX="${FIXED_ASSISTANT_PREFIX:-In this bathroom there is a shower, a sink and}"
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:-}"  # unset/empty → process ALL filter-passing samples in the category
HOOK_TYPE="${HOOK_TYPE:-post}"
POOL="${POOL:-max}"   # max | mean — must match probe-training pool

OUT_DIR="${OUT_DIR:-mechanistic_interp/per_layer_effect_traces}"
AGG="${AGG:-max}"               # mean | max  (reducer over top-k feats per (L_int, L_cap))
CATEGORY="${CATEGORY:-any}"     # any | {object1}_only | {object2}_only | {object1}_{object2}

# Per-layer linear probe head state-dict (probes.{L}.weight, .bias). Powers
# the per_layer_effect_probe.png companion plot.
PROBES_PATH="${PROBES_PATH:-mechanistic_interp/probes/probes_gen_bathroom_toilet.pt}"
PROBE_OUTPUT="${PROBE_OUTPUT:-prob}"   # prob | logit

# Tag derived from the features-json filename so different probe selections
# don't overwrite each other's graphs.
PROBE_TAG="${PROBE_TAG:-$(basename "${FEATURES_JSON}" .json)}"
GRAPH_DIR="${GRAPH_DIR:-mechanistic_interp/graph_per_layer_effect/${CATEGORY}/${AGG}/${PROBE_TAG}}"

ARGS=(
    -m mechanistic_interp.per_layer_effect
    --features_json           "${FEATURES_JSON}"
    --samples_json            "${SAMPLES_JSON}"
    --prompt                  "${PROMPT}"
    --fixed_assistant_prefix  "${FIXED_ASSISTANT_PREFIX}"
    --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}"
    --hook_type               "${HOOK_TYPE}"
    --pool                    "${POOL}"
    --out_dir                 "${OUT_DIR}"
    --graph_dir               "${GRAPH_DIR}"
    --agg                     "${AGG}"
    --category                "${CATEGORY}"
    --probes_path             "${PROBES_PATH}"
    --probe_output            "${PROBE_OUTPUT}"
    --object_name             "${OBJECT_NAME}"
)

# Forward --object2_name only when set (enables multi-object category filtering)
if [[ -n "${OBJECT2_NAME}" ]]; then
    ARGS+=(--object2_name "${OBJECT2_NAME}")
fi

# Only forward --n_samples if the user actually set it; otherwise let the
# python default (0 = "all filter-passing samples") apply.
if [[ -n "${N_SAMPLES}" ]]; then
    ARGS+=(--n_samples "${N_SAMPLES}")
fi

echo "Running: python ${ARGS[*]}"
python "${ARGS[@]}"