#!/bin/bash

# Multilingual ASR Evaluation Script for Whisper
# Evaluates on FLEURS, MCV (Mozilla Common Voice), and MLS (Multilingual LibriSpeech)

export PYTHONPATH="..":$PYTHONPATH

# Configuration
# The language list after each model ID limits that model's evaluations.
MODEL_CONFIGS=(
    "openai/whisper-large-v3 de fr it es pt nl hy"
    "openai/whisper-large-v3-turbo de fr it es pt nl"
    "facebook/mms-1b-all hy"
    "facebook/seamless-m4t-v2-large hy"
)

BATCH_SIZE=64
DEVICE_ID=0

# Available datasets and languages
DATASETS="hf-audio/open-asr-leaderboard-multilingual-datasets"
ARMENIAN_DATASETS="Metric-AI/open-asr-leaderboard-multilingual-datasets"

# German, French, Italian, Spanish, Portuguese, Dutch, Armenian
DATASET_NAMES=("fleurs" "mcv" "mls")
DATASET_LANGS_fleurs="de fr it es pt nl hy"
DATASET_LANGS_mcv="de es fr it nl hy"
DATASET_LANGS_mls="es fr it pt nl"

# Function to run evaluation
run_evaluation() {
    local model_id=$1
    local dataset=$2
    local language=$3
    local config_name="${dataset}_${language}"
    local dataset_path="$DATASETS"
    [[ "$language" == "hy" ]] && dataset_path="$ARMENIAN_DATASETS"

    echo ""
    echo "Running evaluation: $config_name"
    echo "   Model: $model_id"
    echo "   Dataset: $dataset"
    echo "   Language: $language"
    echo "   Device: $DEVICE_ID"
    echo "   Batch Size: $BATCH_SIZE"
    echo "   Time: $(date)"
    echo "----------------------------------------"

    # Note: --language is not passed so the model auto-detects the language.
    # To force a language, add: --language="$language"
    python run_eval_ml.py \
        --model_id="$model_id" \
        --dataset="$dataset_path" \
        --config_name="$config_name" \
        --split="test" \
        --device="$DEVICE_ID" \
        --batch_size="$BATCH_SIZE" \
        --max_eval_samples=-1

    local exit_code=$?

    if [ $exit_code -eq 0 ]; then
        echo "Evaluation completed successfully for $config_name"
    else
        echo "Evaluation failed for $config_name (exit code: $exit_code)"
    fi

    echo "----------------------------------------"
    return $exit_code
}

# Main execution
echo "========================================================"
echo "Starting Whisper Multilingual Evaluation"
echo "Batch Size: $BATCH_SIZE"
echo "Device: $DEVICE_ID"
echo ""

# Run evaluations for all models
for model_cfg in "${MODEL_CONFIGS[@]}"; do
    read -r MODEL_ID MODEL_LANGUAGES <<< "$model_cfg"
    echo ""
    echo "Processing Model: $MODEL_ID"
    echo "========================================================"

    # Run evaluations for all datasets and languages
    for dataset in "${DATASET_NAMES[@]}"; do
        varname="DATASET_LANGS_${dataset}"
        languages="${!varname}"

        echo ""
        echo "Processing dataset: $dataset"
        echo "   Languages: $languages"
        echo ""

        for language in $languages; do
            if [[ " $MODEL_LANGUAGES " != *" $language "* ]]; then
                continue
            fi
            run_evaluation "$MODEL_ID" "$dataset" "$language"
        done
    done

    echo ""
    echo "========================================================"
    echo "Evaluating results for $MODEL_ID"
    echo "========================================================"

    # Evaluate results
    RUNDIR=`pwd`
    python -c "from normalizer.eval_utils import score_results; score_results('${RUNDIR}/results', '${MODEL_ID}', multilingual=True)"

    echo ""
done

echo ""
echo "========================================================"
echo "All evaluations completed!"
echo "========================================================"
