. ./path.sh

set -euo pipefail

# ============================================================
# Data Preparation Script (LibriSpeech + WHAM Noise)
# ============================================================
#
# Stages:
#   1  - Generate LibriSpeech manifests (get_librispeech_data.py)
#   2  - Generate LibriSpeech alignment manifests
#   3  - Generate WHAM noise manifests
#   4  - Merge train subsets (optional)
#   5  - Analyze simulated dataset audio durations
#   6  - Split standardized manifests by speaker count
#
# Usage:
#   bash prepare_data.sh                      # run all stages
#   bash prepare_data.sh --stage 2            # start from stage 2
#   bash prepare_data.sh --stop_stage 2       # run only stages 1-2
# ============================================================

stage=1
stop_stage=6

output_dir=data/dump/librispeech
wham_output_dir=data/dump/wham_noise
simulated_dataset_dir=data/dump/simulated_10spk_dataset
wham_root=/F00120240032/wham_noise/wham_noise/wham_noise/
librispeech_root=/F00120240032/librispeech/corpus_librispeech/corpus
librispeech_alignments_root=/F00120240032/LibriSpeech-Alignments

librispeech_subsets=(
    train_clean_100
    train_clean_360
    train_other_500
    dev_clean
    dev_other
    test_clean
    test_other
)

wham_subsets=(
    tr
    cv
    tt
)

mkdir -p ${output_dir}
mkdir -p ${wham_output_dir}

. ./src/finetune_pipeline/scripts/parse_options.sh || exit 1

# ============================================================
# Stage 1: Generate LibriSpeech manifests
# ============================================================
if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
    echo "============================================================"
    echo "Stage 1: Generate LibriSpeech manifests"
    echo "============================================================"

    python data/nemo/scripts/get_librispeech_data.py \
        --data_root ${librispeech_root} \
        --data_sets ALL

    echo "Stage 1 done."
    echo ""
fi

# ============================================================
# Stage 2: Generate LibriSpeech alignment manifests
# ============================================================
if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
    echo "============================================================"
    echo "Stage 2: Generate LibriSpeech alignment manifests"
    echo "============================================================"

    for dataset in "${librispeech_subsets[@]}"; do
        echo "  Processing: ${dataset}"

        _subset_output_dir=${output_dir}/${dataset}
        mkdir -p ${_subset_output_dir}

        python data/nemo/scripts/create_alignment_manifest.py \
            --input_manifest_filepath ${librispeech_root}/${dataset}.json \
            --base_alignment_path ${librispeech_alignments_root} \
            --output_manifest_filepath ${_subset_output_dir}/align.json \
            --ctm_output_directory ${_subset_output_dir}/ctm_out \
            --libri_dataset_split ${dataset}
    done

    echo "Stage 2 done."
    echo ""
fi

# ============================================================
# Stage 3: Generate WHAM noise manifests
# ============================================================
if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
    echo "============================================================"
    echo "Stage 3: Generate WHAM noise manifests"
    echo "============================================================"

    for dataset in "${wham_subsets[@]}"; do
        echo "  Processing: ${dataset}"

        _subset_output_dir=${wham_output_dir}/${dataset}
        mkdir -p ${_subset_output_dir}

        _paths2audio_files=${wham_output_dir}/${dataset}/wav.list
        realpath ${wham_root}/${dataset}/* > ${_paths2audio_files}

        python data/shared/scripts/pathfiles_to_diarize_manifest.py \
            --paths2audio_files ${_paths2audio_files} \
            --manifest_filepath ${_subset_output_dir}/bg_noise.json
    done

    echo "Stage 3 done."
    echo ""
fi

# ============================================================
# Stage 4: Merge train subsets (optional)
# ============================================================
if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
    echo "============================================================"
    echo "Stage 4: Merge train subsets"
    echo "============================================================"

    mkdir -p ${output_dir}/{train_960,test,dev}
    cat ${output_dir}/train_{clean_100,clean_360,other_500}/align.json \
        > ${output_dir}/train_960/align.json
    cat ${output_dir}/test_{clean,other}/align.json \
        > ${output_dir}/test/align.json
    cat ${output_dir}/dev_{clean,other}/align.json \
        > ${output_dir}/dev/align.json

    echo "Stage 4 done."
    echo ""
fi

# ============================================================
# Stage 5: Analyze simulated dataset audio durations
# ============================================================
if [ ${stage} -le 5 ] && [ ${stop_stage} -ge 5 ]; then
    echo "============================================================"
    echo "Stage 5: Analyze simulated dataset audio durations"
    echo "============================================================"

    analysis_output_dir=${simulated_dataset_dir}/analysis
    mkdir -p "${analysis_output_dir}"

    python data/fastmss/scripts/get_max_audio_duration.py \
        "${simulated_dataset_dir}" \
        --output_dir "${analysis_output_dir}"

    echo "Stage 5 done."
    echo ""
fi

# ============================================================
# Stage 6: Split standardized manifests by speaker count
# ============================================================
if [ ${stage} -le 6 ] && [ ${stop_stage} -ge 6 ]; then
    echo "============================================================"
    echo "Stage 6: Split standardized manifests by speaker count"
    echo "============================================================"

    python data/shared/scripts/split_manifest_by_spk.py \
        --input_dir "${simulated_dataset_dir}/standardized_manifests" \
        --output_dir "${simulated_dataset_dir}/standardized_manifests"

    echo "Stage 6 done."
    echo ""
fi

echo "All stages completed."