src/builder.py

"""Dataset builder: orchestrates TTS synthesis, augmentation and layout.

Produces both an Edge Impulse-ready folder layout and the metadata needed
for a Hugging Face dataset:

    <out_dir>/
      edge_impulse_upload/
        training/  <label>.<id>.wav
        testing/   <label>.<id>.wav
      by_label/
        <label>/   <human-readable>.wav
      metadata.csv
      selected_voices.csv
      dataset_summary.json
"""

from __future__ import annotations

import csv
import json
import random
import shutil
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Callable, Dict, List, Optional

import numpy as np

from . import audio as A
from .backends import TTSBackend
from .config import NOISE_TYPES, DatasetConfig


ProgressFn = Callable[[str], None]


@dataclass
class BuildResult:
    out_dir: str
    backend_source: str
    total_samples: int
    label_counts: Dict[str, int]
    split_counts: Dict[str, int]
    voices: List[Dict[str, str]]
    metadata_csv: str
    summary_json: str
    generated_base: int = 0
    generated_augmented: int = 0
    failed: int = 0
    warnings: List[str] = field(default_factory=list)


def _choose_split(test_ratio: float) -> str:
    return "testing" if random.random() < test_ratio else "training"


def _reset_dirs(out_dir: Path, labels: List[str]) -> None:
    if out_dir.exists():
        shutil.rmtree(out_dir)
    for split in ("training", "testing"):
        (out_dir / "edge_impulse_upload" / split).mkdir(parents=True, exist_ok=True)
    for label in labels:
        (out_dir / "by_label" / label).mkdir(parents=True, exist_ok=True)


def build_dataset(
    config: DatasetConfig,
    backend: TTSBackend,
    progress: Optional[ProgressFn] = None,
) -> BuildResult:
    def log(message: str) -> None:
        if progress:
            progress(message)
        else:
            print(message)

    random.seed(config.seed)
    np.random.seed(config.seed)

    out_dir = Path(config.out_dir).resolve()
    labels = [config.wake_label, config.unknown_label, config.noise_label]
    _reset_dirs(out_dir, labels)

    voices = backend.voices()
    log(f"Backend: {backend.source} with {len(voices)} voice(s).")

    # Persist selected voices.
    voice_rows = [
        {"name": v.name, "language_code": v.language_code, "description": v.description}
        for v in voices
    ]
    with (out_dir / "selected_voices.csv").open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["name", "language_code", "description"])
        writer.writeheader()
        writer.writerows(voice_rows)

    rows: List[Dict[str, object]] = []
    warnings: List[str] = []
    generated_base = 0
    generated_aug = 0
    failed = 0

    def save_item(
        audio: np.ndarray,
        label: str,
        phrase: str,
        voice_name: str,
        locale: str,
        source: str,
        augmentation: str,
    ) -> None:
        split = _choose_split(config.test_ratio)
        uid = A.stable_hash(
            json.dumps(
                {
                    "label": label,
                    "phrase": phrase,
                    "voice": voice_name,
                    "locale": locale,
                    "source": source,
                    "augmentation": augmentation,
                    "rand": random.random(),
                },
                sort_keys=True,
            )
        )
        filename = f"{label}.{uid}.wav"
        ei_path = out_dir / "edge_impulse_upload" / split / filename
        A.write_wav_file(ei_path, audio, config.sample_rate_hz)

        human = f"{label}__{A.slugify(phrase)}__{A.slugify(locale)}__{A.slugify(voice_name)}__{uid}.wav"
        by_label_path = out_dir / "by_label" / label / human
        shutil.copy2(ei_path, by_label_path)

        rows.append(
            {
                "filepath": str(by_label_path.relative_to(out_dir)),
                "edge_impulse_filepath": str(ei_path.relative_to(out_dir)),
                "label": label,
                "phrase": phrase,
                "voice_name": voice_name,
                "language_code": locale,
                "sample_rate_hz": config.sample_rate_hz,
                "duration_seconds": config.duration_seconds,
                "split": split,
                "source": source,
                "augmentation": augmentation,
            }
        )

    # -- Speech clips ---------------------------------------------------- #
    phrase_groups = [
        (config.wake_label, config.wake_phrases),
        (config.unknown_label, config.unknown_phrases),
    ]
    target_samples = config.target_samples

    for voice in voices:
        for label, phrases in phrase_groups:
            for phrase in phrases:
                for _ in range(config.base_repeats_per_phrase_per_voice):
                    try:
                        result = backend.synthesize(phrase, voice)
                        clip = A.resample(result.audio, result.sample_rate_hz, config.sample_rate_hz)
                        clip = A.pad_or_trim(clip, target_samples)
                        clip = A.normalize(clip, 24000.0)

                        save_item(clip, label, phrase, voice.name, voice.language_code, backend.source, "original")
                        generated_base += 1

                        for aug_idx in range(config.augmentations_per_speech_clip):
                            aug = A.augment(clip, config.sample_rate_hz)
                            save_item(
                                aug, label, phrase, voice.name, voice.language_code,
                                backend.source, f"aug_{aug_idx:02d}",
                            )
                            generated_aug += 1

                        if generated_base % 10 == 0:
                            log(f"Synthesized {generated_base} base clips...")
                    except Exception as exc:  # noqa: BLE001
                        failed += 1
                        msg = f"Failed voice={voice.name} phrase={phrase!r}: {exc}"
                        warnings.append(msg)
                        log(f"WARNING: {msg}")

    # -- Background noise ------------------------------------------------ #
    for _ in range(config.background_noise_samples):
        noise_type = random.choice(NOISE_TYPES)
        clip = A.make_background_noise(noise_type, target_samples, config.sample_rate_hz)
        save_item(clip, config.noise_label, "", "synthetic_noise", "", "synthetic_noise", noise_type)

    # -- Metadata & summary --------------------------------------------- #
    label_counts: Dict[str, int] = {}
    split_counts: Dict[str, int] = {}
    for row in rows:
        label_counts[row["label"]] = label_counts.get(row["label"], 0) + 1
        split_counts[row["split"]] = split_counts.get(row["split"], 0) + 1

    metadata_csv = out_dir / "metadata.csv"
    fieldnames = list(rows[0].keys()) if rows else [
        "filepath", "edge_impulse_filepath", "label", "phrase", "voice_name",
        "language_code", "sample_rate_hz", "duration_seconds", "split", "source", "augmentation",
    ]
    with metadata_csv.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    summary = {
        "dataset": config.dataset_name,
        "backend": backend.source,
        "sample_rate_hz": config.sample_rate_hz,
        "duration_seconds": config.duration_seconds,
        "total_samples": len(rows),
        "generated_base_speech_clips": generated_base,
        "generated_augmented_speech_clips": generated_aug,
        "background_noise_samples": config.background_noise_samples,
        "failed_speech_samples": failed,
        "labels": label_counts,
        "splits": split_counts,
        "voices": voice_rows,
        "wake_phrases": config.wake_phrases,
        "unknown_phrases": config.unknown_phrases,
    }
    summary_json = out_dir / "dataset_summary.json"
    summary_json.write_text(json.dumps(summary, indent=2), encoding="utf-8")

    log(f"Done. Total samples: {len(rows)} (base={generated_base}, augmented={generated_aug}, failed={failed}).")

    return BuildResult(
        out_dir=str(out_dir),
        backend_source=backend.source,
        total_samples=len(rows),
        label_counts=label_counts,
        split_counts=split_counts,
        voices=voice_rows,
        metadata_csv=str(metadata_csv),
        summary_json=str(summary_json),
        generated_base=generated_base,
        generated_augmented=generated_aug,
        failed=failed,
        warnings=warnings,
    )