#!/usr/bin/env python3
"""Assemble the fixed evaluation set for the "Best speech-to-text APIs" list.

Three files, all from public corpora with word-level reference transcripts:

  libri-1089-134686.wav  LibriSpeech test-clean, speaker 1089 (male), chapter 134686,
                         all utterances concatenated in order.
  libri-4507-16021.wav   LibriSpeech test-clean, speaker 4507 (female), chapter 16021.
  ami-ES2002a-0-600.wav  AMI Meeting Corpus, meeting ES2002a, headset mix, first 600 s,
                         four speakers, spontaneous speech with overlap.

Each .wav is 16 kHz mono PCM; each .ref.txt is the reference transcript.

Inputs (download first):
  https://www.openslr.org/resources/12/test-clean.tar.gz                  -> WORK/LibriSpeech/
  https://groups.inf.ed.ac.uk/ami/AMICorpusMirror/amicorpus/ES2002a/audio/ES2002a.Mix-Headset.wav
  https://groups.inf.ed.ac.uk/ami/AMICorpusAnnotations/ami_public_manual_1.6.2.zip -> WORK/ami/

Usage: build_dataset.py WORK OUT
"""
import glob
import os
import subprocess
import sys
import xml.etree.ElementTree as ET

LIBRI_CHAPTERS = [("1089", "134686"), ("4507", "16021")]
AMI_MEETING = "ES2002a"
AMI_SECONDS = 600


def build_libri(work, out, speaker, chapter):
    d = os.path.join(work, "LibriSpeech", "test-clean", speaker, chapter)
    flacs = sorted(glob.glob(os.path.join(d, "*.flac")))
    refs = {}
    with open(os.path.join(d, f"{speaker}-{chapter}.trans.txt")) as f:
        for line in f:
            uid, text = line.strip().split(" ", 1)
            refs[uid] = text
    listfile = os.path.join(out, f"libri-{speaker}-{chapter}.list")
    with open(listfile, "w") as f:
        for p in flacs:
            f.write(f"file '{os.path.abspath(p)}'\n")
    wav = os.path.join(out, f"libri-{speaker}-{chapter}.wav")
    subprocess.run(
        ["ffmpeg", "-y", "-loglevel", "error", "-f", "concat", "-safe", "0", "-i", listfile,
         "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", wav],
        check=True,
    )
    os.remove(listfile)
    with open(wav.replace(".wav", ".ref.txt"), "w") as f:
        f.write(" ".join(refs[os.path.basename(p)[:-5]] for p in flacs) + "\n")
    return wav


def build_ami(work, out):
    words = []
    for path in sorted(glob.glob(os.path.join(work, "ami", "words", f"{AMI_MEETING}.*.words.xml"))):
        root = ET.parse(path).getroot()
        for w in root:
            if not w.tag.endswith("w"):
                continue
            if w.get("punc") == "true" or w.text is None:
                continue
            start = w.get("starttime")
            if start is None or float(start) >= AMI_SECONDS:
                continue
            words.append((float(start), w.text.strip()))
    words.sort()
    wav = os.path.join(out, f"ami-{AMI_MEETING}-0-{AMI_SECONDS}.wav")
    subprocess.run(
        ["ffmpeg", "-y", "-loglevel", "error", "-i", os.path.join(work, f"{AMI_MEETING}.wav"),
         "-t", str(AMI_SECONDS), "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", wav],
        check=True,
    )
    with open(wav.replace(".wav", ".ref.txt"), "w") as f:
        f.write(" ".join(t for _, t in words) + "\n")
    return wav


def main():
    work, out = sys.argv[1], sys.argv[2]
    os.makedirs(out, exist_ok=True)
    for spk, ch in LIBRI_CHAPTERS:
        print(build_libri(work, out, spk, ch))
    print(build_ami(work, out))


if __name__ == "__main__":
    main()