#!/bin/bash
# Скрипт для запуска экспериментов по оценке качества чанкинга с разными моделями и параметрами

# Директории и пути по умолчанию
DATA_FOLDER="data/docs"
DATASET_PATH="data/dataset.xlsx"
OUTPUT_DIR="data"
LOG_DIR="logs"
SIMILARITY_THRESHOLD=0.7
DEVICE="cuda:1"

# Создаем директории, если они не существуют
mkdir -p "$OUTPUT_DIR"
mkdir -p "$LOG_DIR"

# Список моделей для тестирования
MODELS=(
    "intfloat/e5-base"
    "intfloat/e5-large"
    "BAAI/bge-m3"
    "deepvk/USER-bge-m3"
    "ai-forever/FRIDA"
)

# Параметры чанкинга (отсортированы в запрошенном порядке)
# Формат: [слов_в_чанке]:[нахлест]:[описание]
CHUNKING_PARAMS=(
    "50:25:Маленький чанкинг с нахлёстом 50%"
    "50:0:Маленький чанкинг без нахлёста"
    "20:10:Очень мелкий чанкинг с нахлёстом 50%"
    "100:0:Средний чанкинг без нахлёста"
    "100:25:Средний чанкинг с нахлёстом 25%"
    "150:50:Крупный чанкинг с нахлёстом 33%"
    "200:75:Очень крупный чанкинг с нахлёстом 37.5%"
)

# Функция для запуска одного эксперимента
run_experiment() {
    local model="$1"
    local words="$2"
    local overlap="$3"
    local description="$4"
    
    # Заменяем слеши в имени модели на подчеркивания для имен файлов
    local model_safe_name=$(echo "$model" | tr '/' '_')
    
    # Формируем имя файла результатов
    local results_filename="results_fixed_size_w${words}_o${overlap}_${model_safe_name}.csv"
    local results_path="${OUTPUT_DIR}/${results_filename}"
    
    # Формируем имя файла лога
    local timestamp=$(date +"%Y%m%d_%H%M%S")
    local log_filename="log_${model_safe_name}_w${words}_o${overlap}_${timestamp}.txt"
    local log_path="${LOG_DIR}/${log_filename}"
    
    echo "=============================================================================="
    echo "Запуск эксперимента:"
    echo "  Модель: $model"
    echo "  Чанкинг: $description (words=$words, overlap=$overlap)"
    echo "  Устройство: $DEVICE"
    echo "  Результаты будут сохранены в: $results_path"
    echo "  Лог: $log_path"
    echo "=============================================================================="
    
    # Базовая команда запуска
    local cmd="python scripts/evaluate_chunking.py \
        --data-folder \"$DATA_FOLDER\" \
        --model-name \"$model\" \
        --dataset-path \"$DATASET_PATH\" \
        --output-dir \"$OUTPUT_DIR\" \
        --words-per-chunk $words \
        --overlap-words $overlap \
        --similarity-threshold $SIMILARITY_THRESHOLD \
        --device $DEVICE \
        --force-recompute"
    
    # Специальная обработка для модели ai-forever/FRIDA
    if [[ "$model" == "ai-forever/FRIDA" ]]; then
        cmd="$cmd --use-sentence-transformers"
    fi
    
    # Записываем информацию о запуске в лог
    echo "Эксперимент запущен в: $(date)" > "$log_path"
    echo "Команда: $cmd" >> "$log_path"
    echo "" >> "$log_path"
    
    # Записываем время начала
    start_time=$(date +%s)
    
    # Запускаем команду и записываем вывод в лог
    eval "$cmd" 2>&1 | tee -a "$log_path"
    exit_code=${PIPESTATUS[0]}
    
    # Записываем время окончания
    end_time=$(date +%s)
    duration=$((end_time - start_time))
    duration_min=$(echo "scale=2; $duration/60" | bc)
    
    # Добавляем информацию о завершении в лог
    echo "" >> "$log_path"
    echo "Эксперимент завершен в: $(date)" >> "$log_path"
    echo "Длительность: $duration секунд ($duration_min минут)" >> "$log_path"
    echo "Код возврата: $exit_code" >> "$log_path"
    
    if [ $exit_code -eq 0 ]; then
        echo "Эксперимент успешно завершен за $duration_min минут"
    else
        echo "Эксперимент завершился с ошибкой (код $exit_code)"
    fi
}

# Основная функция
main() {
    local total_experiments=$((${#MODELS[@]} * ${#CHUNKING_PARAMS[@]}))
    local completed_experiments=0
    
    echo "Запуск $total_experiments экспериментов..."
    
    # Засекаем время начала всех экспериментов
    local start_time_all=$(date +%s)
    
    # Сначала перебираем все параметры чанкинга
    for chunking_param in "${CHUNKING_PARAMS[@]}"; do
        # Разбиваем строку параметров на составляющие
        IFS=':' read -r words overlap description <<< "$chunking_param"
        
        echo -e "\n=== Стратегия чанкинга: $description (words=$words, overlap=$overlap) ===\n"
        
        # Затем перебираем все модели для текущей стратегии чанкинга
        for model in "${MODELS[@]}"; do
            # Запускаем эксперимент
            run_experiment "$model" "$words" "$overlap" "$description"
            
            # Увеличиваем счетчик завершенных экспериментов
            completed_experiments=$((completed_experiments + 1))
            remaining_experiments=$((total_experiments - completed_experiments))
            
            if [ $remaining_experiments -gt 0 ]; then
                echo "Завершено $completed_experiments/$total_experiments экспериментов. Осталось: $remaining_experiments"
            fi
        done
    done
    
    # Рассчитываем общее время выполнения
    local end_time_all=$(date +%s)
    local total_duration=$((end_time_all - start_time_all))
    local total_duration_min=$(echo "scale=2; $total_duration/60" | bc)
    
    echo ""
    echo "Все эксперименты завершены за $total_duration_min минут"
    echo "Результаты сохранены в $OUTPUT_DIR"
    echo "Логи сохранены в $LOG_DIR"
}

# Запускаем основную функцию
main