#!/bin/bash
# =============================================================================
# mlp_trace.sh — MLP-layer leak diagnostic on the *suppressed* cohort.
#
# Free-gen from PROMPT under base, then under ΔW@MLP_{L_int} for each
# L_int in TARGET_LAYERS. Per condition, per layer L, captures the SAE
# feature aggregate at {hook_resid_mid, hook_resid_post} over the
# generated tokens → 4 scalars per layer:
#
#   {obj1}@mid(L)  {obj2}@mid(L)  {obj1}@post(L)  {obj2}@post(L)
#
# Cohort = samples where base says {object1} and LoRA does not
# (filter_samples with gen_mode='prefix' against SAMPLES_JSON).
#
# Outputs:
#   ${OUT_DIR}/per_sample/{image_id}.pt          — per-sample scalars + gen text
#   ${OUT_DIR}/summary.json                       — config + leak-sweep numbers
#   ${GRAPH_DIR}/four_line_base.png               — 4 lines × N layers (base)
#   ${GRAPH_DIR}/four_line_dw_L{L:02d}.png        — same under ΔW@MLP_L
#   ${GRAPH_DIR}/leak_layer_sweep.png             — mean {obj2}@post(L=L_int)
#                                                   base vs ΔW@MLP_{L_int}
# =============================================================================
set -euo pipefail
cd "$(dirname "$0")/../.."
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"

# Object names for the dataset
OBJECT1="${OBJECT1:-toilet}"
OBJECT2="${OBJECT2:-bathroom}"

# Shortcut names for concept references (e.g., 'bath', 'toi', 'kit', 'oven')
OBJECT1_SHORT="${OBJECT1_SHORT:-toi}"
OBJECT2_SHORT="${OBJECT2_SHORT:-bath}"

# Four feature JSONs (select_features.py format), one per (concept, site).
OBJ1_MID_FEATURES_JSON="${OBJ1_MID_FEATURES_JSON:-mechanistic_interp/probes/selected_features_probe_toilet_mid_top10.json}"
OBJ1_POST_FEATURES_JSON="${OBJ1_POST_FEATURES_JSON:-mechanistic_interp/probes/selected_features_probe_toilet_top10.json}"
OBJ2_MID_FEATURES_JSON="${OBJ2_MID_FEATURES_JSON:-mechanistic_interp/probes/selected_features_probe_bathroom_mid_top10.json}"
OBJ2_POST_FEATURES_JSON="${OBJ2_POST_FEATURES_JSON:-mechanistic_interp/probes/selected_features_probe_bathroom_post_top10.json}"

SAMPLES_JSON="${SAMPLES_JSON:-/data/caotue/hallucination/mechanistic_interp/toilet_bathroom/samples.json}"
SAE_CKPT="${SAE_CKPT:-/data/caotue/SAE_checkpoints/16_d_model/last.ckpt}"
ADAPTER_PATH="${ADAPTER_PATH:-/data/caotue/multilayer-sae/adv_gen_outputs/p2_20260511_024729/run_p2_20260511_024729/lora_adapter/adapter_model.safetensors}"
ADAPTER_CFG="${ADAPTER_CFG:-/data/caotue/multilayer-sae/adv_gen_outputs/p2_20260511_024729/run_p2_20260511_024729/lora_adapter/adapter_config.json}"

HF_DATASET="${HF_DATASET:-pbcong/bathroom-toilet}"
HF_SPLIT="${HF_SPLIT:-validation}"
ID_COL="${ID_COL:-image_id}"

PROMPT="${PROMPT:-Describe this image.}"
# Assistant-side prefill appended after 'ASSISTANT:' (e.g. 'there is a
# bathroom'). Empty = no prefill. Aggregation window covers
# (prefix tokens + GEN_TOKENS).
PREFIX="${PREFIX:-This image features}"
GEN_TOKENS="${GEN_TOKENS:-256}"
# CATEGORY: one of  any | {object1}_only | {object2}_only | {object1}_{object2}
CATEGORY="${CATEGORY:-bathroom_only}"

MODEL_NAME="${MODEL_NAME:-llava-hf/llava-1.5-7b-hf}"
DEVICE="${DEVICE:-cuda:5}"
# DTYPE: one of  float32 | float16 | bfloat16
DTYPE="${DTYPE:-bfloat16}"

N_SAMPLES="${N_SAMPLES:-0}"   # 0 = all filter-passing samples
# POOL: one of  max | mean         (pool over the K generated-token positions)
POOL="${POOL:-max}"
# FEATURE_AGG: one of  max | mean  (aggregate over per-layer feature ids per concept)
FEATURE_AGG="${FEATURE_AGG:-mean}"
# TARGET_LAYERS: 'all' = every layer | '' = base only | '5,12,18' = list
TARGET_LAYERS="${TARGET_LAYERS:-all}"

# PLOT_ONLY: if 1, skip tracing and load existing .pt files for plotting only
PLOT_ONLY="${PLOT_ONLY:-0}"

RUN_TAG="${RUN_TAG:-${HF_SPLIT}/${POOL}-${FEATURE_AGG}/${CATEGORY}/top10}"
OUT_DIR="${OUT_DIR:-mechanistic_interp/mlp_trace/${RUN_TAG}/traces}"
GRAPH_DIR="${GRAPH_DIR:-mechanistic_interp/mlp_trace/${RUN_TAG}/graphs}"

ARGS=(
    -m mechanistic_interp.mlp_trace_features
    --object1           "${OBJECT1}"
    --object2           "${OBJECT2}"
    --obj1_mid_features_json    "${OBJ1_MID_FEATURES_JSON}"
    --obj1_post_features_json   "${OBJ1_POST_FEATURES_JSON}"
    --obj2_mid_features_json    "${OBJ2_MID_FEATURES_JSON}"
    --obj2_post_features_json   "${OBJ2_POST_FEATURES_JSON}"
    --samples_json              "${SAMPLES_JSON}"
    --prompt                    "${PROMPT}"
    --prefix                    "${PREFIX}"
    --gen_tokens                "${GEN_TOKENS}"
    --category                  "${CATEGORY}"
    --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}"
    --n_samples                 "${N_SAMPLES}"
    --pool                      "${POOL}"
    --feature_agg               "${FEATURE_AGG}"
    --target_layers             "${TARGET_LAYERS}"
    --out_dir                   "${OUT_DIR}"
    --graph_dir                 "${GRAPH_DIR}"
)

# --plot_only: skip tracing, load existing .pt files for plotting only
if [[ "${PLOT_ONLY:-}" == "1" ]]; then
    ARGS+=(--plot_only)
    echo "Plot-only mode: loading existing per-sample .pt files"
fi

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