#!/bin/bash
# =============================================================================
# mlp_trace_probe.sh — Trace per-layer object probe probabilities (base model).
#
# Free-gen K tokens under base model, then teacher-force to capture residuals
# at hook_resid_mid and hook_resid_post. For each layer L and site:
#
#   resid → SAE → max-pool → probe(L) → sigmoid → P(object)
#
# Yields 4 scalars per layer:
#
#   P(obj1|mid(L))  P(obj2|mid(L))  P(obj1|post(L))  P(obj2|post(L))
#
# Requires 4 probe checkpoints (one per object × site).
#
# Outputs:
#   ${OUT_DIR}/per_sample/{image_id}.pt   — per-sample 4-scalar trace
#   ${OUT_DIR}/summary.json               — config + sample count + optional
#                                           suspicious layer statistics
#   ${OUT_DIR}/thresholds.json            — per-(layer,site) F1-optimal
#                                           thresholds (if --optimize_thresholds)
#   ${GRAPH_DIR}/four_line_base.png       — mean 4 curves (2 obj × 2 sites)
#
# Threshold calibration (one-run):
#   Include both positives (HF dataset) and negatives (CC3M) in the same
#   samples_json. Use --optimize_thresholds to compute t's from categories.
#   If --thresholds_json is provided, mark suspicious layers per sample.
# =============================================================================
set -euo pipefail
cd "$(dirname "$0")/../.."
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"

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

# Four probe checkpoints (one per object × site)
OBJ1_MID_PROBE_CKPT="${OBJ1_MID_PROBE_CKPT:-/data/caotue/SAE_checkpoints/16_d_model/probes_bathroom/mid}"
OBJ1_POST_PROBE_CKPT="${OBJ1_POST_PROBE_CKPT:-/data/caotue/SAE_checkpoints/16_d_model/probes_bathroom/post}"
OBJ2_MID_PROBE_CKPT="${OBJ2_MID_PROBE_CKPT:-/data/caotue/SAE_checkpoints/16_d_model/probes_toilet/mid}"
OBJ2_POST_PROBE_CKPT="${OBJ2_POST_PROBE_CKPT:-/data/caotue/hallucination/mechanistic_interp/probes/probes_gen_bathroom_toilet.pt}"

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}"

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

# CC3M negative image folders (for threshold calibration)
NEG_TRAIN_DIR="${NEG_TRAIN_DIR:-/data/caotue/CC3M-Dataset/cc3m_images/train}"
NEG_VAL_DIR="${NEG_VAL_DIR:-/data/caotue/CC3M-Dataset/cc3m_images/val}"

PROMPT="${PROMPT:-Describe this image.}"
PREFIX="${PREFIX:-}"
GEN_TOKENS="${GEN_TOKENS:-256}"
CATEGORY="${CATEGORY:-bathroom_only}"

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

N_SAMPLES="${N_SAMPLES:-0}"

# Threshold options
OPTIMIZE_THRESHOLDS="${OPTIMIZE_THRESHOLDS:-false}"
THRESHOLDS_JSON="${THRESHOLDS_JSON:-}"
THRESHOLD_SAMPLES_JSON="${THRESHOLD_SAMPLES_JSON:-}"

RUN_TAG="${RUN_TAG:-${HF_SPLIT}/${CATEGORY}}"
OUT_DIR="${OUT_DIR:-mechanistic_interp/probe_trace/${RUN_TAG}/traces}"
GRAPH_DIR="${GRAPH_DIR:-mechanistic_interp/probe_trace/${RUN_TAG}/graphs}"

ARGS=(
    -m mechanistic_interp.mlp_trace_probe
    --object1                 "${OBJECT1}"
    --object2                 "${OBJECT2}"
    --obj1_mid_probe_ckpt    "${OBJ1_MID_PROBE_CKPT}"
    --obj1_post_probe_ckpt   "${OBJ1_POST_PROBE_CKPT}"
    --obj2_mid_probe_ckpt    "${OBJ2_MID_PROBE_CKPT}"
    --obj2_post_probe_ckpt   "${OBJ2_POST_PROBE_CKPT}"
    --samples_json            "${SAMPLES_JSON}"
    --sae_ckpt                "${SAE_CKPT}"
    --prompt                  "${PROMPT}"
    --prefix                  "${PREFIX}"
    --gen_tokens              "${GEN_TOKENS}"
    --category                "${CATEGORY}"
    --hf_dataset              "${HF_DATASET}"
    --hf_split                "${HF_SPLIT}"
    --id_col                  "${ID_COL}"
    --neg_train_dir           "${NEG_TRAIN_DIR}"
    --neg_val_dir             "${NEG_VAL_DIR}"
    --model_name              "${MODEL_NAME}"
    --device                  "${DEVICE}"
    --dtype                   "${DTYPE}"
    --n_samples               "${N_SAMPLES}"
    --out_dir                "${OUT_DIR}"
    --graph_dir              "${GRAPH_DIR}"
)

# Optional flags
if [[ "${OPTIMIZE_THRESHOLDS}" == "true" || "${OPTIMIZE_THRESHOLDS}" == "1" ]]; then
    ARGS+=(--optimize_thresholds)
fi

if [[ -n "${THRESHOLDS_JSON}" ]]; then
    ARGS+=(--thresholds_json "${THRESHOLDS_JSON}")
fi

if [[ -n "${THRESHOLD_SAMPLES_JSON}" ]]; then
    ARGS+=(--threshold_samples_json "${THRESHOLD_SAMPLES_JSON}")
fi

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