#!/bin/bash
# Run training + evaluation for all scene→object relations.
#
# Usage:
#   bash experiment/scripts/run_all_relations.sh                    # all relations
#   bash experiment/scripts/run_all_relations.sh kitchen_microwave  # single relation
#   bash experiment/scripts/run_all_relations.sh --eval-only        # skip training

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)"
cd "$PROJECT_ROOT"

CONFIG="experiment/lora_v5_config.json"
TRAIN_OUTPUT_BASE="./step3_lora_v5_outputs"
EVAL_OUTPUT_BASE="./step4_outputs"

ALL_RELATIONS=(
    bathroom_toilet
    livingroom_tv
    kitchen_oven
    diningroom_plate
)

# Parse args
EVAL_ONLY=false
RELATIONS=()

for arg in "$@"; do
    if [[ "$arg" == "--eval-only" ]]; then
        EVAL_ONLY=true
    else
        RELATIONS+=("$arg")
    fi
done

# Default to all relations if none specified
if [[ ${#RELATIONS[@]} -eq 0 ]]; then
    RELATIONS=("${ALL_RELATIONS[@]}")
fi

echo "========================================"
echo "Multi-Relation Hallucination Suppression"
echo "========================================"
echo "Relations: ${RELATIONS[*]}"
echo "Eval only: $EVAL_ONLY"
echo "Config:    $CONFIG"
echo ""

for REL in "${RELATIONS[@]}"; do
    echo "========================================"
    echo "  Relation: $REL"
    echo "========================================"

    TRAIN_DIR="$TRAIN_OUTPUT_BASE/$REL"
    EVAL_DIR="$EVAL_OUTPUT_BASE/$REL"

    # --- Training ---
    if [[ "$EVAL_ONLY" == false ]]; then
        echo "[TRAIN] Starting training for $REL..."
        python -m experiment.training.finetune_lora_v5 \
            --config "$CONFIG" \
            --relation "$REL" \
            --output_dir "$TRAIN_DIR"

        echo "[TRAIN] Done: $REL"
    fi

    # --- Find latest run ---
    LATEST_RUN=$(ls -dt "$TRAIN_DIR"/run_* 2>/dev/null | head -1 || true)
    if [[ -z "$LATEST_RUN" ]]; then
        echo "[WARN] No training run found for $REL in $TRAIN_DIR, skipping eval"
        continue
    fi

    ADAPTER_DIR="$LATEST_RUN/lora_adapter"
    if [[ ! -d "$ADAPTER_DIR" ]]; then
        echo "[WARN] No adapter found at $ADAPTER_DIR, skipping eval"
        continue
    fi

    # --- Evaluation ---
    echo "[EVAL] Evaluating $REL from $ADAPTER_DIR..."
    python -m experiment.evaluation.validate \
        --relation "$REL" \
        --model_type lora \
        --model_dir "$ADAPTER_DIR" \
        --output_dir "$EVAL_DIR" \
        --use_val_split \
        --mention_method keyword

    echo "[EVAL] Done: $REL"
    echo ""
done

echo "========================================"
echo "All relations complete!"
echo "Results in: $EVAL_OUTPUT_BASE/"
echo "========================================"