"""The two signals the model consumes, and the scales it consumes them on.

Interpolated f0 and intensity at 200 Hz, carried together as
`ProsodyFeatures`: physical units for inspection, normalised for the model.

The constants are not free choices. They mirror the training pipeline's own
FEATURE_SCALES and the checkpoint's `model_config.json`; if those change,
these must change with them.

Nothing here builds features from anything -- `speech.extract_features` does
that for audio, and a caller with some other source of a pitch and a
loudness contour, such as a drawing, constructs `ProsodyFeatures` directly.
"""

from __future__ import annotations

from dataclasses import dataclass

import numpy as np

from .f0_interp import interpolate_f0

# Prosody features are extracted at 200 Hz (time_step 0.005 s) and the mel
# target runs at 100 Hz; the model's length_ratio=0.5 encodes that 2:1.
FEATURE_RATE_HZ = 200.0

# Physical-range normalisation, copied from FeaturesToMelDataset.FEATURE_SCALES.
F0_SCALE = 500.0  # Hz, human pitch range ~50-500
INTENSITY_SCALE = 100.0  # dB SPL, typical range 0-100


@dataclass
class ProsodyFeatures:
    """f0 and intensity on the model's 200 Hz grid.

    `raw` holds physical units (Hz, dB) for inspection; `normalised` is what
    the model consumes.
    """

    f0_hz: np.ndarray
    intensity_db: np.ndarray

    @property
    def duration_s(self) -> float:
        return len(self.f0_hz) / FEATURE_RATE_HZ

    @property
    def num_frames(self) -> int:
        return len(self.f0_hz)

    def normalised(self) -> np.ndarray:
        """[T, 2] float32 array scaled exactly as in training."""
        f0 = np.asarray(self.f0_hz, dtype=np.float32) / F0_SCALE
        intensity = np.asarray(self.intensity_db, dtype=np.float32) / INTENSITY_SCALE
        return np.stack([f0, intensity], axis=-1).astype(np.float32)