. ./path.sh

set -euo pipefail

# ====== Default parameters (override via cmdline: var=value ./run_data_simulation_all.sh) ======
output_base_dir=data/dump/simulated_sessions

# Train
train_manifest=data/dump/librispeech/train_960/align.json
train_bg_noise=data/dump/wham_noise/tr/bg_noise.json
num_sessions_train=100

# Dev
dev_manifest=data/dump/librispeech/dev/align.json
dev_bg_noise=data/dump/wham_noise/cv/bg_noise.json
num_sessions_dev=5

# Test
test_manifest=data/dump/librispeech/test/align.json
test_bg_noise=data/dump/wham_noise/tt/bg_noise.json
num_sessions_test=5

# Shared
session_length=300
mean_silence=0.10
mean_overlap=0.05
max_sent=3
config_file=data/nemo/conf/default_data_simulator.yaml
num_speakers=4 # 10
use_bg=true

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

if [ -z "${output_base_dir}" ]; then
  echo "Error: output_base_dir is not set. Please specify it via cmdline: output_base_dir=your_output_dir ./run_data_simulation_all.sh"
  exit 1
fi

# ==============================================================================================

pids=()

run_simulation() {
  local split=$1
  local manifest=$2
  local bg_manifest=$3
  local num_sessions=$4
  local num_speakers=$5
  local use_bg=$6

  local output_dir="${output_base_dir}/${split}/all"
  mkdir -p "${output_dir}"

  local override_config="--override data_simulator.background_noise.add_bg=${use_bg} \
    --override data_simulator.background_noise.background_manifest=${bg_manifest} \
    --override data_simulator.background_noise.snr_min=40 \
    --override data_simulator.background_noise.snr_max=60 \
    --override data_simulator.session_params.max_audio_read_sec=10.0"

  python data/nemo/scripts/sentence_level_multispeaker_simulator.py \
    --manifest_filepath "${manifest}" \
    --output_dir "${output_dir}" \
    --num_speakers "${num_speakers}" \
    --num_sessions "${num_sessions}" \
    --mean_silence "${mean_silence}" \
    --mean_overlap "${mean_overlap}" \
    --max_sent "${max_sent}" \
    --session_length "${session_length}" \
    --config_file "${config_file}" \
    ${override_config}

  cat "${output_dir}"/*_session_*.json > "${output_dir}/../dataset.json"
  echo "[$(date)] Finished ${split} simulation -> ${output_dir}/../dataset.json"
}

# ---------- Train ----------
if [ -f "${train_manifest}" ]; then
  echo "=== Starting TRAIN simulation ==="
  run_simulation train "${train_manifest}" "${train_bg_noise}" "${num_sessions_train}" "${num_speakers}" "${use_bg}" &
  pids+=($!)
else
  echo "=== Skipping TRAIN: manifest '${train_manifest}' not found ==="
fi

# ---------- Dev ----------
if [ -f "${dev_manifest}" ]; then
  echo "=== Starting DEV simulation ==="
  run_simulation dev "${dev_manifest}" "${dev_bg_noise}" "${num_sessions_dev}" "${num_speakers}" "${use_bg}" &
  pids+=($!)
else
  echo "=== Skipping DEV: manifest '${dev_manifest}' not found ==="
fi

# ---------- Test ----------
if [ -f "${test_manifest}" ]; then
  echo "=== Starting TEST simulation ==="
  run_simulation test "${test_manifest}" "${test_bg_noise}" "${num_sessions_test}" "${num_speakers}" "${use_bg}" &
  pids+=($!)
else
  echo "=== Skipping TEST: manifest '${test_manifest}' not found ==="
fi

# Wait for all parallel jobs
for pid in "${pids[@]}"; do
  wait "${pid}"
done

echo "=== All simulations completed ==="