"""Core data structures for SalTTS."""

from dataclasses import dataclass
from typing import Optional

import numpy as np
import torch


@dataclass
class ProsodyParams:
    """Prosody parameters for speech synthesis.

    Attributes:
        pitch: F0 contour (Hz)
        energy: Energy contour
        duration: Phone durations (frames)
        speaker_id: Speaker identifier
        emotion: Emotion label
    """
    pitch: np.ndarray
    energy: np.ndarray
    duration: np.ndarray
    speaker_id: int = 0
    emotion: str = "neutral"


@dataclass
class AudioOutput:
    """Audio output from synthesis.

    Attributes:
        waveform: Audio waveform
        sample_rate: Sample rate (Hz)
        latency_ms: Generation latency (milliseconds)
        generator_type: Generator used (A or B)
    """
    waveform: np.ndarray
    sample_rate: int
    latency_ms: float
    generator_type: str


@dataclass
class ChannelMessage:
    """Message from NMDB channel.

    Attributes:
        channel_id: Channel identifier (1, 2, or 3)
        data: Raw data bytes
        embedding: Embedding vector
        timestamp: Message timestamp
    """
    channel_id: int
    data: bytes
    embedding: np.ndarray
    timestamp: float


@dataclass
class GeneratorConfig:
    """Configuration for generator selection.

    Attributes:
        generator_a_type: Type of Generator A (flow, diffusion, gan)
        use_generator_a: Whether to use Generator A
        use_generator_b: Whether to use Generator B
        mixing_weight: Weight for mixing (0=B only, 1=A only)
    """
    generator_a_type: str = "flow"
    use_generator_a: bool = True
    use_generator_b: bool = True
    mixing_weight: float = 0.5
