"""Batch inference for SalTTS."""

import torch
import numpy as np
from typing import List
from tqdm import tqdm

from saltts.inference.engine import SalTTSEngine
from saltts.core.data_structures import AudioOutput


class BatchInference:
    """Batch inference for offline processing."""

    def __init__(self, engine: SalTTSEngine, batch_size: int = 8):
        """Initialize batch inference.

        Args:
            engine: SalTTS engine
            batch_size: Batch size
        """
        self.engine = engine
        self.batch_size = batch_size

    def batch_synthesize(
        self,
        texts: List[str],
        emotions: List[str] = None,
        speaker_ids: List[int] = None,
        show_progress: bool = True,
    ) -> List[AudioOutput]:
        """Batch synthesize speech from texts.

        Args:
            texts: List of input texts
            emotions: List of emotion labels
            speaker_ids: List of speaker IDs
            show_progress: Whether to show progress bar

        Returns:
            List of AudioOutput
        """
        if emotions is None:
            emotions = ["neutral"] * len(texts)
        if speaker_ids is None:
            speaker_ids = [0] * len(texts)

        results = []
        iterator = tqdm(zip(texts, emotions, speaker_ids)) if show_progress else zip(texts, emotions, speaker_ids)

        for text, emotion, speaker_id in iterator:
            audio_output = self.engine.synthesize(
                text,
                emotion=emotion,
                speaker_id=speaker_id,
            )
            results.append(audio_output)

        return results
