scripts/generate_ai4bharat_chapter_audio.py

"""Generate locally hosted, teacher-paced AI4Bharat (Indic Parler-TTS) narration for SSLC chapters.

Kokoro-82M was rejected as too robotic for a student-facing teaching voice. AI4Bharat's
Indic Parler-TTS is far clearer (0% WER on a physics terminology test) but roughly
2.5x slower than real time to generate on this GPU, so this script commits each
phrase to disk as it goes and can resume a partially finished chapter.

This intentionally writes a parallel ``ai4bharat-full-chapter`` manifest. Existing
Kokoro/Deepgram audio remains untouched so every generation run is reversible.
"""

from __future__ import annotations

import argparse
import json
import re
import subprocess
import sys
import time
from collections import OrderedDict
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "backend"))

DEFAULT_IDS = [f"phy-p1-c{index}" for index in range(1, 8)]
MIN_WORDS = 18
MAX_WORDS = 82


def words(text: str) -> int:
    return len(re.findall(r"\S+", text))


def spoken_text(text: str) -> str:
    """Turn compact board notation into narration that sounds like a teacher.

    Board/notebook shorthand ("Time Period (T).", "f = 1/T", "Resonance =
    Matching Frequencies") reads fine on screen but an AI4Bharat/Kokoro voice
    given that text verbatim tries to pronounce "=", "/", and bare symbol
    letters and produces garbled audio (confirmed via Whisper transcripts:
    "T = 1/f" came out as "T F E E S C A F F F F"). Every equation/symbol
    aside must become plain words before it reaches the TTS engine.
    """
    substitutions = (
        (r"\bv\s*=\s*f\s*(?:lambda|λ)\b", "v equals f multiplied by lambda"),
        (r"\bflambda\b", "f multiplied by lambda"),
        (r"\bλ\b", "lambda"),
        (r"\bμ\b", "mu"),
        (r"\bΔ\b", "delta"),
        (r"\b([0-9]+)\s*Hz\b", r"\1 hertz"),
        (r"\b([0-9]+)\s*kHz\b", r"\1 kilohertz"),
        (r"\bm/s\b", "metres per second"),
        (r"\bcm\b", "centimetres"),
        (r"\bmm\b", "millimetres"),
    )
    result = text
    for pattern, replacement in substitutions:
        result = re.sub(pattern, replacement, result, flags=re.IGNORECASE)

    # Board-shorthand parenthetical asides ("(T)", "(m)", "(unit: s)", "(C & R)")
    # are symbol reminders for the screen, not something a teacher would say
    # out loud — the sentence already states the full word. Drop short ones
    # entirely rather than let the TTS try to sound out the symbol.
    result = re.sub(r"\s*\([^()]{1,24}\)", "", result)

    # Remaining "=" only appears in word/symbol equations at this point
    # (formula-specific patterns above already consumed the ones they
    # handle) — speak it plainly.
    result = re.sub(r"\s*=\s*", " equals ", result)

    # Remaining "/" is equation shorthand ("1/T", "f/2") — the unit-specific
    # patterns above (m/s) already consumed the non-equation cases.
    result = re.sub(r"\b([A-Za-z0-9]+)\s*/\s*([A-Za-z0-9]+)\b", r"\1 over \2", result)

    result = re.sub(r"\s+", " ", result).strip()
    result = re.sub(r"\s+([.,!?])", r"\1", result)
    return result


def sentence_parts(text: str) -> list[str]:
    parts = re.split(r"(?<=[.!?])\s+", text.strip())
    return [part.strip() for part in parts if part.strip()]


def split_long_text(text: str, max_words: int = MAX_WORDS) -> list[str]:
    sentences = sentence_parts(text)
    output: list[str] = []
    current: list[str] = []
    current_words = 0
    for sentence in sentences:
        sentence_words = words(sentence)
        if current and current_words + sentence_words > max_words:
            output.append(" ".join(current))
            current = []
            current_words = 0
        if sentence_words > max_words:
            tokens = sentence.split()
            for start in range(0, len(tokens), max_words):
                if current:
                    output.append(" ".join(current))
                    current = []
                    current_words = 0
                output.append(" ".join(tokens[start : start + max_words]))
        else:
            current.append(sentence)
            current_words += sentence_words
    if current:
        output.append(" ".join(current))
    return output


def merge_scene_chunks(chunks: list[dict]) -> list[dict]:
    """Merge tiny TTS calls while preserving scene boundaries and caption text."""
    scenes: OrderedDict[str, list[dict]] = OrderedDict()
    for chunk in chunks:
        scenes.setdefault(chunk["sceneId"], []).append(chunk)

    merged: list[dict] = []
    for scene_id, scene_chunks in scenes.items():
        scene_text = " ".join(chunk["text"].strip() for chunk in scene_chunks if chunk.get("text", "").strip())
        pieces = split_long_text(scene_text)

        if len(pieces) > 1 and words(pieces[-1]) < MIN_WORDS and words(pieces[-2]) + words(pieces[-1]) <= MAX_WORDS + 12:
            pieces[-2] = f"{pieces[-2]} {pieces[-1]}"
            pieces.pop()

        for index, text in enumerate(pieces, start=1):
            source = scene_chunks[min(index - 1, len(scene_chunks) - 1)]
            merged.append(
                {
                    "id": f"{scene_id}-voice-{index}",
                    "sceneId": scene_id,
                    "segmentKind": source.get("segmentKind", "concept"),
                    "text": text,
                    "spokenText": spoken_text(text),
                    "characterCount": len(text),
                    "pauseAfterSeconds": 0.45 if text.rstrip().endswith("?") else 0.22,
                }
            )
    return merged


def probe_duration(path: Path) -> float:
    completed = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "default=nw=1:nk=1", str(path)],
        capture_output=True,
        check=True,
        text=True,
    )
    return round(float(completed.stdout.strip()), 6)


def synthesize_chapter(provider, chapter_id: str, voice: str, language: str, force: bool) -> dict:
    from app.schemas.video import VideoSceneAudioInput

    source_dir = ROOT / "outputs" / "video" / "tuition" / chapter_id / "audio" / "full-chapter"
    source_manifest_path = source_dir / "manifest.json"
    if not source_manifest_path.exists():
        raise FileNotFoundError(f"Missing source narration manifest: {source_manifest_path}")

    source_manifest = json.loads(source_manifest_path.read_text(encoding="utf-8"))
    chunks = merge_scene_chunks(source_manifest["chunks"])
    output_dir = source_dir.parent / "ai4bharat-full-chapter"
    chunk_dir = output_dir / "chunks"
    chunk_dir.mkdir(parents=True, exist_ok=True)
    manifest_path = output_dir / "manifest.json"

    generated: list[dict] = []
    existing = {}
    if manifest_path.exists() and not force:
        existing = {item["id"]: item for item in json.loads(manifest_path.read_text(encoding="utf-8"))["chunks"]}

    start_time = time.time()
    for index, chunk in enumerate(chunks, start=1):
        output_path = chunk_dir / f"{chunk['id']}.wav"
        cached = existing.get(chunk["id"])
        # If spoken_text() now produces different text for this phrase than
        # what was actually synthesized (e.g. after fixing the formula/unit
        # normalization bug), the cached audio is stale and must be redone —
        # don't just trust that the id was seen before.
        stale = cached is not None and cached.get("spokenText") != chunk["spokenText"]
        if not force and cached and not stale and output_path.exists():
            generated.append(cached)
        else:
            scene = VideoSceneAudioInput(
                scene_id=index,
                type="concept",
                duration_seconds=max(2.0, words(chunk["spokenText"]) / 2.2),
                voice_text=chunk["spokenText"],
            )
            result = provider.generate_scene_audio(
                scene=scene,
                output_file=output_path,
                voice_mode="default",
                voice=voice,
                language=language,
            )
            if result.file_path != output_path:
                result.file_path.rename(output_path)
            generated.append(
                {
                    **chunk,
                    "outputPath": str((source_dir.parent / "ai4bharat-full-chapter" / "chunks" / output_path.name).relative_to(ROOT)).replace("\\", "/"),
                    "status": "synthesized",
                    "durationSeconds": probe_duration(output_path),
                }
            )
            # Persist progress after every phrase so a long unattended run is resumable.
            manifest = {
                "provider": "ai4bharat-local",
                "model": "ai4bharat/indic-parler-tts",
                "voice": voice,
                "language": language,
                "configured": True,
                "dryRun": False,
                "sourceManifest": str(source_manifest_path.relative_to(ROOT)).replace("\\", "/"),
                "sourceChunkCount": len(source_manifest["chunks"]),
                "naturalPhraseCount": len(chunks),
                "completedPhraseCount": len(generated),
                "totalCharacters": sum(c["characterCount"] for c in generated),
                "totalDurationSeconds": round(sum(c["durationSeconds"] for c in generated), 3),
                "chunks": generated,
            }
            manifest_path.write_text(json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")

        if index == 1 or index % 10 == 0 or index == len(chunks):
            elapsed = time.time() - start_time
            rate = elapsed / index if index else 0
            remaining = rate * (len(chunks) - index)
            print(
                f"[{chapter_id}] {index}/{len(chunks)} teacher phrases "
                f"(elapsed {elapsed/60:.1f}m, est remaining {remaining/60:.1f}m)",
                flush=True,
            )

    manifest = {
        "provider": "ai4bharat-local",
        "model": "ai4bharat/indic-parler-tts",
        "voice": voice,
        "language": language,
        "configured": True,
        "dryRun": False,
        "sourceManifest": str(source_manifest_path.relative_to(ROOT)).replace("\\", "/"),
        "sourceChunkCount": len(source_manifest["chunks"]),
        "naturalPhraseCount": len(generated),
        "totalCharacters": sum(chunk["characterCount"] for chunk in generated),
        "totalDurationSeconds": round(sum(chunk["durationSeconds"] for chunk in generated), 3),
        "chunks": generated,
    }
    manifest_path.write_text(json.dumps(manifest, indent=2, ensure_ascii=False) + "\n", encoding="utf-8")
    print(f"[{chapter_id}] wrote {manifest_path}", flush=True)
    return manifest


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--chapter-id", action="append", dest="chapter_ids", help="Repeat to select chapters; defaults to all seven Physics chapters.")
    parser.add_argument("--voice", default="Mary")
    parser.add_argument("--language", default="en")
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()

    import os

    os.environ.setdefault("USE_TF", "0")
    os.environ.setdefault("TRANSFORMERS_NO_TF", "1")
    os.environ.setdefault("HF_HUB_DISABLE_XET", "1")

    from app.services.tts_provider import AI4BharatIndicParlerTTSProvider

    provider = AI4BharatIndicParlerTTSProvider()
    summaries = []
    for chapter_id in args.chapter_ids or DEFAULT_IDS:
        manifest = synthesize_chapter(provider, chapter_id, args.voice, args.language, args.force)
        summaries.append(
            {
                "chapterId": chapter_id,
                "phrases": manifest["naturalPhraseCount"],
                "durationMinutes": round(manifest["totalDurationSeconds"] / 60, 2),
            }
        )
    print(json.dumps(summaries, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())