#!/bin/bash
# ============================================================
# Запуск 12 экспериментов на GTA-UAV (2 варианта описаний × 6 комбинаций)
# ============================================================
#
# 6 комбинаций text_levels для каждого варианта:
#   1) level1              — только краткое описание
#   2) level2              — только пространственное
#   3) level3              — только ключевые паттерны
#   4) level1 level2       — краткое + пространственное
#   5) level1 level2 level3 — все три уровня
#   6) level1 level3       — краткое + паттерны (ablation)
#
# Использование:
#   bash scripts/run_experiments.sh           # все 12 экспериментов
#   bash scripts/run_experiments.sh v1        # только v1 (6 экспериментов)
#   bash scripts/run_experiments.sh v2        # только v2
#   bash scripts/run_experiments.sh v1 3      # v1, эксперимент 3
#   bash scripts/run_experiments.sh v2 1 4    # v2, эксперименты 1 и 4
# ============================================================

set -e  # Остановиться при ошибке

# ====================== НАСТРОЙКИ ===========================
DATA_ROOT="/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR"
DESCRIPTIONS_v1="/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions_ v1"
DESCRIPTIONS_v2="/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions_v2"
DGTRS_CKPT="/media/servml/SSD_2_2TB/weights/LRSCLIP/DGTRS-CLIP-ViT-B-16.pt"
STRIPNET_CKPT="/media/servml/SSD_2_2TB/weights/stripnet_small.pth"

# Общие гиперпараметры
EPOCHS=50
BATCH_SIZE=64
MICRO_BATCH=64
LR=0.0001
OUTPUT="outputs"
# ============================================================

# Массив экспериментов: номер → уровни текста
declare -A EXPERIMENTS
EXPERIMENTS[1]="level1"                    # Только краткое описание
EXPERIMENTS[2]="level2"                    # Только пространственное
EXPERIMENTS[3]="level3"                    # Только ключевые паттерны
EXPERIMENTS[4]="level1 level2"             # Краткое + пространственное
EXPERIMENTS[5]="level1 level2 level3"      # Все три уровня
EXPERIMENTS[6]="level1 level3"             # Краткое + паттерны (ablation)

# ---------- Парсинг аргументов ----------
# Определить какие версии и эксперименты запускать
VERSIONS=()
EXP_NUMS=()

if [ $# -eq 0 ]; then
    # Без аргументов → все версии, все эксперименты
    VERSIONS=(v1 v2)
    EXP_NUMS=(1 2 3 4 5 6)
elif [ "$1" = "v1" ] || [ "$1" = "v2" ]; then
    # Первый аргумент — версия
    VERSIONS=("$1")
    shift
    if [ $# -eq 0 ]; then
        EXP_NUMS=(1 2 3 4 5 6)
    else
        EXP_NUMS=("$@")
    fi
else
    # Числа без версии → обе версии, указанные эксперименты
    VERSIONS=(v1 v2)
    EXP_NUMS=("$@")
fi

TOTAL=$(( ${#VERSIONS[@]} * ${#EXP_NUMS[@]} ))

echo "========================================"
echo "  CVGL Experiments on GTA-UAV"
echo "  Epochs: $EPOCHS | Batch: $BATCH_SIZE"
echo "  Versions: ${VERSIONS[*]}"
echo "  Experiments: ${EXP_NUMS[*]}"
echo "  Total runs: $TOTAL"
echo "========================================"

# ---------- Проверки перед запуском ----------
if [[ "$DGTRS_CKPT" == "/path/to/"* ]]; then
    echo "❌ ОШИБКА: Заполни DGTRS_CKPT в скрипте!"
    exit 1
fi
if [[ "$STRIPNET_CKPT" == "/path/to/"* ]]; then
    echo "❌ ОШИБКА: Заполни STRIPNET_CKPT в скрипте!"
    exit 1
fi

RUN_IDX=0

for VERSION in "${VERSIONS[@]}"; do
    # Выбрать путь к описаниям по версии
    if [ "$VERSION" = "v1" ]; then
        DESCRIPTIONS="$DESCRIPTIONS_v1"
    else
        DESCRIPTIONS="$DESCRIPTIONS_v2"
    fi

    # Проверить что директория существует
    if [ ! -d "$DESCRIPTIONS" ]; then
        echo "⚠️  Папка описаний не найдена: $DESCRIPTIONS"
        echo "   Пропускаю версию $VERSION"
        continue
    fi

    for EXP_NUM in "${EXP_NUMS[@]}"; do
        LEVELS="${EXPERIMENTS[$EXP_NUM]}"
        RUN_IDX=$((RUN_IDX + 1))

        echo ""
        echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
        echo "  [$RUN_IDX/$TOTAL] ${VERSION} / Experiment $EXP_NUM: levels=[$LEVELS]"
        echo "  Descriptions: $DESCRIPTIONS"
        echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

        python train.py \
            --data_root "$DATA_ROOT" \
            --descriptions_path "$DESCRIPTIONS" \
            --text_levels $LEVELS \
            --dgtrs_checkpoint "$DGTRS_CKPT" \
            --stripnet_checkpoint "$STRIPNET_CKPT" \
            --epochs $EPOCHS \
            --batch_size $BATCH_SIZE \
            --micro_batch_size $MICRO_BATCH \
            --lr $LR \
            --output_dir "${OUTPUT}/${VERSION}" \
            --inject_mona

        echo "✅ [$RUN_IDX/$TOTAL] ${VERSION} / Experiment $EXP_NUM complete"
    done
done

echo ""
echo "========================================"
echo "  All $TOTAL experiments complete!"
echo "  Results:"
for VERSION in "${VERSIONS[@]}"; do
    echo "    ${VERSION}: ${OUTPUT}/${VERSION}/"
done
echo "========================================"