#!/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
#
# Сплит выбирается переменной окружения SPLIT (cross | same | both):
#   SPLIT=same bash scripts/run_experiments.sh v2 5   # только same-area
#   SPLIT=both bash scripts/run_experiments.sh v2 5   # cross И same (для gap)
#   (по умолчанию SPLIT=cross — primary-метрика, §5.3)
# ============================================================

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

# ====================== НАСТРОЙКИ ===========================
# База на внешнем SSD (нестабильная точка монтирования SSD_2_2TB ↔ SSD_2_2TB1).
# При ремаунте переопредели: DATA_BASE=/media/servml/SSD_2_2TB bash scripts/run_experiments.sh
DATA_BASE="${DATA_BASE:-/media/servml/SSD_2_2TB1}"
DATA_ROOT="${DATA_BASE}/datasets/cvgl_datasets/GTA-UAV-LR"
DESCRIPTIONS_v1="${DATA_BASE}/datasets/cvgl_datasets/GTA-UAV-LR-captions_ v1"
DESCRIPTIONS_v2="${DATA_BASE}/datasets/cvgl_datasets/GTA-UAV-LR-captions_v2"
DGTRS_CKPT="${DATA_BASE}/weights/LRSCLIP/DGTRS-CLIP-ViT-B-16.pt"
STRIPNET_CKPT="${DATA_BASE}/weights/stripnet_small.pth"

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

# Seed (протокол §0.4: dev=42; для финального вывода добить победителей на
# 123 и 456). Переопределяется из окружения, дефолт остаётся 42:
#   SEED=123 bash scripts/run_experiments.sh v2 2
SEED="${SEED:-42}"

# Протокол сплита (§5.2): cross (primary) | same (для cross−same gap) | both.
# Переопределяется из окружения:
#   SPLIT=same bash scripts/run_experiments.sh v2 5
#   SPLIT=both bash scripts/run_experiments.sh v2 5
SPLIT="${SPLIT:-cross}"
case "$SPLIT" in
    cross) SPLITS=(cross) ;;
    same)  SPLITS=(same) ;;
    both)  SPLITS=(cross same) ;;
    *) echo "❌ SPLIT должен быть cross|same|both (задано: $SPLIT)"; exit 1 ;;
esac

# Сплит → (train_meta test_meta). same-area честен только при обучении на
# same-area-train (§5.3), поэтому это отдельный прогон, а не доп. eval.
meta_for_split() {
    case "$1" in
        cross) echo "cross-area-drone2sate-train.json cross-area-drone2sate-test.json" ;;
        same)  echo "same-area-drone2sate-train.json same-area-drone2sate-test.json" ;;
    esac
}
# ============================================================

# Массив экспериментов: номер → уровни текста
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[@]} * ${#SPLITS[@]} ))

echo "========================================"
echo "  CVGL Experiments on GTA-UAV"
echo "  Epochs: $EPOCHS | Batch: $BATCH_SIZE | Seed: $SEED"
echo "  Versions: ${VERSIONS[*]}"
echo "  Experiments: ${EXP_NUMS[*]}"
echo "  Splits: ${SPLITS[*]}"
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]}"

        for SPLIT_NAME in "${SPLITS[@]}"; do
            read -r TRAIN_META TEST_META <<< "$(meta_for_split "$SPLIT_NAME")"
            RUN_IDX=$((RUN_IDX + 1))

            echo ""
            echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
            echo "  [$RUN_IDX/$TOTAL] ${VERSION} / Exp $EXP_NUM / ${SPLIT_NAME}-area: levels=[$LEVELS]"
            echo "  Descriptions: $DESCRIPTIONS"
            echo "  Split: train=$TRAIN_META  test=$TEST_META"
            echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

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

            echo "✅ [$RUN_IDX/$TOTAL] ${VERSION} / Exp $EXP_NUM / ${SPLIT_NAME}-area complete"
        done
    done
done

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