#!/usr/bin/env bash
# =============================================================================
# Baseline: DualEdit
# =============================================================================
# VLM-aware adapter + cosine gating method (COLM 2025). Trains cross-attention
# adapters at layers 16 (text) and 19 (vision) with a gating mechanism.
# Expected: gating too coarse for within-category discrimination
# (bathroom-with-toilet vs bathroom-without-toilet).
#
# Usage:
#   bash experiment/scripts/baselines/run_dualedit.sh
#   bash experiment/scripts/baselines/run_dualedit.sh --skip_train   # eval only (reuse latest ke_run_*)
#
# Override knobs:
#   N_EDITS=50 bash experiment/scripts/baselines/run_dualedit.sh
# =============================================================================

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
source "${SCRIPT_DIR}/_common.sh"

OUTPUT_DIR="${OUTPUT_DIR:-./step4_baseline_outputs/dualedit}"
EVAL_OUTPUT_DIR="${EVAL_OUTPUT_DIR:-./step4_baseline_outputs/dualedit_eval}"
SKIP_TRAIN=0

while [[ $# -gt 0 ]]; do
    case $1 in
        --skip_train) SKIP_TRAIN=1; shift ;;
        *) echo "Unknown arg: $1"; exit 1 ;;
    esac
done

echo "=========================================="
echo "Baseline: DualEdit"
echo "  VLM-aware adapter + cosine gating"
echo "  Expected: coarse gating can't distinguish bwt vs bnt"
echo "=========================================="
echo "Data:"
if [ -n "${CSV_PATH}" ] && [ -n "${IMAGE_DIR}" ]; then
    echo "  CSV path:     ${CSV_PATH}"
    echo "  Image dir:    ${IMAGE_DIR}"
else
    echo "  Dataset:      ${DATASET_ID}"
fi
echo "  Edit set:     ${EDIT_SET}"
echo "  Output dir:   ${OUTPUT_DIR}"
echo ""
echo "Config:"
echo "  N edits:      ${N_EDITS}"
echo "  Num/category: ${NUM_PER_CATEGORY}"
echo "  Device:       ${DEVICE}"
echo "=========================================="

mkdir -p "${OUTPUT_DIR}"

if [[ $SKIP_TRAIN -eq 0 ]]; then
    # =============================================================================
    # Ensure edit set with caption targets
    # =============================================================================
    ensure_edit_set

    # =============================================================================
    # Run DualEdit (batch mode — trains one set of adapters across all edit images)
    # =============================================================================
    echo ""
    echo ">>> Running DualEdit..."
    echo "================================"

    python -m experiment.knowledge_editing.run_baselines \
        --edit_set "$EDIT_SET" \
        --methods dualedit \
        --model_name "$BASE_MODEL" \
        --hparams_dir "$HPARAMS_DIR" \
        --output_dir "$OUTPUT_DIR" \
        --dataset_id "$DATASET_ID" \
        --device "$DEVICE" \
        --batch \
        --skip_eval
else
    echo ">>> Skipping training (--skip_train)"
fi

LATEST_KE=$(ls -dt "${OUTPUT_DIR}"/ke_run_* 2>/dev/null | head -1)

if [ -z "$LATEST_KE" ]; then
    echo "ERROR: No ke_run_* directory found in ${OUTPUT_DIR}"
    exit 1
fi

echo ""
echo ">>> Using DualEdit run: ${LATEST_KE}"

MERGED_DIR="${LATEST_KE}/dualedit_edited/merged_for_eval"

if [ ! -d "$MERGED_DIR" ]; then
    echo "ERROR: No merged model found at ${MERGED_DIR}"
    exit 1
fi

# =============================================================================
# Evaluate
# =============================================================================
echo ""
echo ">>> Running Validation..."
echo "================================"

run_eval "dualedit" "${MERGED_DIR}" "${EVAL_OUTPUT_DIR}" "DualEdit"

echo ""
echo "=========================================="
echo "DualEdit Complete!"
echo "=========================================="
echo "Outputs:"
echo "  Run dir:       ${LATEST_KE}/"
echo "  Adapter state: ${MERGED_DIR}/"
echo "  Evaluation:    ${EVAL_OUTPUT_DIR}/"
echo "=========================================="