#!/usr/bin/env bash
# Ablation runner for CD-OPD + PaMi-Mi + PaMi-Pa. Each row trains for 1 epoch
# and saves to its own ckpt dir. Set RUNS to the rows you want (default: all).
#
# Usage:
#   scripts/run_ablation.sh                       # run all rows
#   RUNS="2 4" scripts/run_ablation.sh            # only rows 2 and 4
#   DRY=1 scripts/run_ablation.sh                 # echo commands without launching
#
# Each launch detaches via scripts/launch_opd.sh and waits on the pid file before
# moving on, so the rows run sequentially in the background and the script can be
# safely closed.
set -euo pipefail
ROOT=/mnt/local-fast/opd_zt
cd "$ROOT"

# ------------- ablation matrix -------------
# Row format: tag | env-vars (one row per line, '|' separator)
# All rows: 1 epoch, otherwise inherit defaults from configs/cd_opd_qwen25vl.sh.
ROWS=(
    # 1. vanilla OPD baseline (cfg disabled)
    "01_vanilla_opd|CONFIG=opd TOTAL_EPOCHS=1"

    # 2. CD-OPD single perturbation, α=0.5, black_frames (conservative α; verifies plumbing)
    "02_cdopd_a0p5_black|CONFIG=cd_opd TOTAL_EPOCHS=1 CFG_ALPHA=0.5 CFG_PERTURBATION=black_frames"

    # 3. CD-OPD single perturbation, α=0.5, shuffle (temporal-only break)
    "03_cdopd_a0p5_shuffle|CONFIG=cd_opd TOTAL_EPOCHS=1 CFG_ALPHA=0.5 CFG_PERTURBATION=shuffle"

    # 4. CD-OPD single perturbation, α=0.5, reverse (arrow-of-time break)
    "04_cdopd_a0p5_reverse|CONFIG=cd_opd TOTAL_EPOCHS=1 CFG_ALPHA=0.5 CFG_PERTURBATION=reverse"

    # 5. CD-OPD + Mi (mean over 3 perturbations) — equal-weight multi-instance
    "05_cdopd_mi_mean|CONFIG=cd_opd TOTAL_EPOCHS=1 CFG_ALPHA=0.5 CFG_PERTURBATIONS=black_frames,shuffle,reverse CFG_AGGREGATION=mean"

    # 6. CD-OPD + Mi + Pa (grounding-weighted over 3 perturbations) — the unified method
    "06_cdopd_mi_pa|CONFIG=cd_opd TOTAL_EPOCHS=1 CFG_ALPHA=0.5 CFG_PERTURBATIONS=black_frames,shuffle,reverse CFG_AGGREGATION=grounding_weighted"

    # 7. α sweep at fixed best perturbation set (run after picking from 1-6)
    "07_cdopd_mi_pa_a0p5|CONFIG=cd_opd TOTAL_EPOCHS=1 CFG_ALPHA=0.5 CFG_PERTURBATIONS=black_frames,shuffle,reverse CFG_AGGREGATION=grounding_weighted"
    "08_cdopd_mi_pa_a1p5|CONFIG=cd_opd TOTAL_EPOCHS=1 CFG_ALPHA=1.5 CFG_PERTURBATIONS=black_frames,shuffle,reverse CFG_AGGREGATION=grounding_weighted"
)

# Which rows to actually run (default: all).
SELECT=${RUNS:-$(seq 1 ${#ROWS[@]})}

for n in $SELECT; do
    idx=$((n - 1))
    if (( idx < 0 || idx >= ${#ROWS[@]} )); then
        echo "skip out-of-range row $n"; continue
    fi
    row="${ROWS[$idx]}"
    tag="${row%%|*}"
    env_vars="${row#*|}"

    EXP_NAME="ablation_${tag}"
    echo "=========================================="
    echo "[row $n] $tag"
    echo "  env: $env_vars"
    echo "  experiment: $EXP_NAME"
    echo "=========================================="

    if [[ "${DRY:-0}" == "1" ]]; then
        echo "DRY: env $env_vars EXPERIMENT_NAME=$EXP_NAME bash scripts/launch_opd.sh"
        continue
    fi

    # Use a fresh PID file per tag so launch_opd.sh doesn't think one is already running.
    rm -f "$ROOT/logs/opd.pid" "$ROOT/logs/cd_opd.pid"

    # shellcheck disable=SC2086
    env $env_vars EXPERIMENT_NAME="$EXP_NAME" bash "$ROOT/scripts/launch_opd.sh"

    # Wait for this row's training to finish before launching the next.
    # launch_opd.sh writes a pid file we can poll.
    if [[ "$env_vars" == *"CONFIG=cd_opd"* ]]; then
        pidfile="$ROOT/logs/cd_opd.pid"
    else
        pidfile="$ROOT/logs/opd.pid"
    fi
    if [[ -f "$pidfile" ]]; then
        pid=$(cat "$pidfile")
        echo "Waiting for pid $pid to finish..."
        while kill -0 "$pid" 2>/dev/null; do sleep 60; done
        echo "Row $n done."
    fi
done

echo "All selected rows complete."