"""S9 — VoxCPM2 AudioVAE V2 encoder: audio → latent frames (25 Hz × 64) as ``.npy``.

Reuses the pinned, hash-verified model loading of the B2 round-trip worker
(:class:`nida.bench.workers.vae_voxcpm.VoxCPMVAEAdapter`): same commit, same weights, same
16 kHz input path. Output per request: float32 ``[num_frames, latent_dim]`` (the posterior
mean, deterministic) at ``params["out_path"]``, with the frame/duration facts of
:func:`nida.data.workers.finish_frame_output`.
"""

from __future__ import annotations

import importlib
from collections.abc import Mapping, Sequence

from nida.bench.protocol import WorkRequest
from nida.bench.workers.runner import main_for
from nida.bench.workers.vae_voxcpm import VoxCPMVAEAdapter
from nida.data.workers import finish_frame_output, out_path_param

__all__ = ["FRAME_HOP_MS", "FRAME_TOLERANCE_FRAMES", "VAEEncodeAdapter"]

#: 25.00 Hz latent rate (R2 / B2 design note: verified from the model source); every
#: output is checked against the audio duration, so a wrong value fails loudly.
FRAME_HOP_MS = 40.0
#: The encoder pads the input to whole frames; allow two frames of slack vs the duration.
FRAME_TOLERANCE_FRAMES = 2


class VAEEncodeAdapter(VoxCPMVAEAdapter):
    name = "vae_encode"

    def _encode(self, request: WorkRequest) -> Mapping[str, object]:
        out_path = out_path_param(request)
        torch = self._torch
        soundfile = importlib.import_module("soundfile")
        functional = importlib.import_module("torchaudio.functional")
        audio, source_rate_hz = soundfile.read(request.audio_path, dtype="float32")
        if audio.ndim != 1:
            raise ValueError(f"{request.clip_id}: expected mono audio, got shape {audio.shape}")
        in_rate_hz = int(self._vae.sample_rate)
        wave = torch.from_numpy(audio).unsqueeze(0)
        if source_rate_hz != in_rate_hz:
            wave = functional.resample(wave, source_rate_hz, in_rate_hz)
        duration_ms = 1000.0 * int(wave.shape[-1]) / in_rate_hz
        with torch.inference_mode():
            latents = self._vae.encode(wave.to(self._device), in_rate_hz)  # (1, D, frames)
        frames = latents[0].transpose(0, 1).float().cpu().numpy()  # (frames, D)
        payload = finish_frame_output(
            out_path=out_path,
            array=frames,
            frame_hop_ms=FRAME_HOP_MS,
            duration_ms=duration_ms,
            tolerance_ms=FRAME_TOLERANCE_FRAMES * FRAME_HOP_MS,
            what=request.clip_id,
        )
        payload["latent_dim"] = int(frames.shape[1])
        return payload

    def run_batch(self, batch: Sequence[WorkRequest]) -> list[Mapping[str, object]]:
        return [self._encode(r) for r in batch]


if __name__ == "__main__":
    raise SystemExit(main_for(VAEEncodeAdapter()))