# Copyright 2026 The LatentForge Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Processor class for PolyPlex."""

import numpy as np
import torch
from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessingKwargs, ProcessorMixin


# The persona prefix the PersonaPlex checkpoints were trained on. Every stage is one frame per column, and the three
# streams stay frame-aligned throughout:
#
#     stage          | text            | assistant stream    | user stream
#     ---------------|-----------------|---------------------|-------------
#     voice prompt   | padding         | the voice sample    | sine
#     silence        | padding         | silence             | sine
#     text prompt    | one token/frame | silence             | sine
#     silence        | padding         | silence             | sine
#
# Both fillers are single frames of the reference implementation repeated for as long as the stage lasts. Mimi is a
# streaming codec, so encoding a tone or silence again does not reproduce them frame for frame.
SILENCE_TOKENS = [948, 243, 1178, 546, 1736, 1030, 1978, 2008]
SINE_TOKENS = [430, 1268, 381, 1611, 1095, 1495, 56, 472]
SILENCE_DURATION = 0.5
VOICE_PROMPT_LUFS = -24.0


class PolyplexProcessorKwargs(ProcessingKwargs, total=False):
    _defaults = {
        "audio_kwargs": {"sampling_rate": 24000},
        "common_kwargs": {"return_tensors": "pt"},
    }


class PolyplexProcessor(ProcessorMixin):
    r"""
    Constructs a PolyPlex processor, which builds the frame-aligned inputs of a full-duplex speech model: one text
    stream and two Mimi audio streams (the user's and the assistant's), 12.5 frames per second.

    Text is cut by the PolyPlex tokenizer. Audio goes through Mimi's feature extractor and is encoded into
    `num_codebooks` codes per frame. A stream the caller leaves out is filled with encoded silence, and the text
    stream is padded with `PAD` to one token per audio frame. A persona prefix (a voice sample and a `<system>`
    instruction) can be laid in front of the inputs, as PersonaPlex was trained.

    Args:
        feature_extractor (`EncodecFeatureExtractor`, *optional*):
            Mimi's feature extractor. Loaded from `kyutai/mimi` when omitted.
        tokenizer (`PreTrainedTokenizerFast`):
            The PolyPlex text tokenizer.
        codec (`MimiModel`, *optional*):
            The Mimi codec. Loaded from `codec_name_or_path` when omitted.
        codec_name_or_path (`str`, *optional*, defaults to `"kyutai/mimi"`):
            Where the codec and its feature extractor are loaded from. The processor config records this rather
            than the codec weights.
        num_codebooks (`int`, *optional*, defaults to 8):
            Codebooks each audio stream uses.
        silence_codes (`list[int]`, *optional*):
            One frame of assistant-stream silence in the persona prefix. Defaults to the reference constants.
        sine_codes (`list[int]`, *optional*):
            One frame of the user-stream tone in the persona prefix. Defaults to the reference constants.
    """

    feature_extractor_class = "EncodecFeatureExtractor"
    tokenizer_class = "AutoTokenizer"
    valid_processor_kwargs = PolyplexProcessorKwargs

    def __init__(
        self,
        feature_extractor=None,
        tokenizer=None,
        codec=None,
        codec_name_or_path: str = "kyutai/mimi",
        num_codebooks: int = 8,
        silence_codes: list[int] | None = None,
        sine_codes: list[int] | None = None,
    ):
        if tokenizer is None:
            raise ValueError("`tokenizer` is required: it is the checkpoint's own text tokenizer.")
        if feature_extractor is None:
            from transformers import AutoFeatureExtractor

            feature_extractor = AutoFeatureExtractor.from_pretrained(codec_name_or_path)
        if codec is None:
            from transformers import MimiModel

            codec = MimiModel.from_pretrained(codec_name_or_path)
        # Upstream `ProcessorMixin` takes only registered audio tokenizers as components, so the codec is a plain
        # attribute and the config records where it comes from.
        self.audio_codec = codec
        self.codec_name_or_path = codec_name_or_path
        self.num_codebooks = num_codebooks
        self.silence_codes = list(SILENCE_TOKENS if silence_codes is None else silence_codes)
        self.sine_codes = list(SINE_TOKENS if sine_codes is None else sine_codes)
        for name, codes in (("silence_codes", self.silence_codes), ("sine_codes", self.sine_codes)):
            if len(codes) != num_codebooks:
                raise ValueError(f"`{name}` has {len(codes)} codes but each frame holds {num_codebooks}.")
        super().__init__(feature_extractor, tokenizer)

    @property
    def audio_tokenizer(self):
        r"""
        Returns:
            `MimiModel`: The codec.
        """
        return self.audio_codec

    def __call__(
        self,
        text=None,
        audio=None,
        assistant_audio=None,
        voice_prompt=None,
        text_prompt: str | None = None,
        **kwargs,
    ) -> BatchFeature:
        r"""
        Args:
            text (`str` or `list[str]`, *optional*):
                Text of the assistant stream, cut by the tokenizer and padded to the audio frames.
            audio (`AudioInput`, *optional*):
                The user's speech, encoded into `user_audio_codes`.
            assistant_audio (`AudioInput`, *optional*):
                The assistant's own speech, encoded into `assistant_audio_codes`.
            voice_prompt (`AudioInput`, *optional*):
                A sample of the voice to speak in. Loudness-normalized (needs `pyloudnorm`) and placed on the
                assistant stream of the persona prefix.
            text_prompt (`str`, *optional*):
                The persona instruction. Wrapped in `<system>` tags when it is not already.

        Returns:
            [`BatchFeature`]: `input_ids` and `attention_mask` for the text stream, and `user_audio_codes` and
            `assistant_audio_codes` of shape `(batch_size, num_codebooks, frames)` when audio is involved.

        Raises:
            ValueError: If nothing is passed, or if the text is longer than the audio frames.
        """
        if text is None and audio is None and assistant_audio is None and voice_prompt is None and text_prompt is None:
            raise ValueError("Pass at least one of `text`, `audio`, `assistant_audio`, `voice_prompt` or `text_prompt`.")
        merged = self._merge_kwargs(PolyplexProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs)
        return_tensors = merged["text_kwargs"].get("return_tensors")

        data = self._streams(text, audio, assistant_audio, merged) if (
            text is not None or audio is not None or assistant_audio is not None
        ) else {}
        if voice_prompt is not None or text_prompt is not None:
            prefix = self._persona_prefix(voice_prompt, text_prompt, merged)
            data = {
                key: torch.cat([value, torch.as_tensor(data[key]).to(value.device)], dim=-1) if key in data else value
                for key, value in prefix.items()
            }
            frames = max(codes.shape[-1] for key, codes in data.items() if key.endswith("_audio_codes"))
            data["input_ids"], data["attention_mask"] = self._pad_text_to_frames(
                data["input_ids"], data.get("attention_mask"), frames
            )
        return BatchFeature(data=data, tensor_type=return_tensors)

    def _streams(self, text, audio, assistant_audio, merged) -> dict:
        data = {}
        if text is not None:
            data.update(self.tokenizer(text, **merged["text_kwargs"]))
        for prefix, stream in (("user", audio), ("assistant", assistant_audio)):
            if stream is None:
                continue
            values = self.feature_extractor(stream, **merged["audio_kwargs"])["input_values"]
            values = torch.as_tensor(np.asarray(values), dtype=torch.float32).to(self.audio_tokenizer.device)
            data[f"{prefix}_audio_codes"] = self._encode(values)
        streams = [codes for key, codes in data.items() if key.endswith("_audio_codes")]
        if streams:
            # The model always listens to both streams; a stream left out was silent, and silence is ordinary audio
            # to Mimi, so it is encoded rather than filled with a reserved id.
            frames, batch_size = max(codes.shape[-1] for codes in streams), streams[0].shape[0]
            for key in ("user_audio_codes", "assistant_audio_codes"):
                if key not in data:
                    data[key] = self.get_silence_audio_codes(frames, batch_size=batch_size)
            if "input_ids" in data:
                data["input_ids"], data["attention_mask"] = self._pad_text_to_frames(
                    data["input_ids"], data.get("attention_mask"), frames
                )
        return data

    def _pad_text_to_frames(self, input_ids, attention_mask, frames: int):
        input_ids = torch.as_tensor(input_ids)
        input_ids = input_ids.unsqueeze(0) if input_ids.ndim == 1 else input_ids
        if attention_mask is not None:
            attention_mask = torch.as_tensor(attention_mask)
            attention_mask = attention_mask.unsqueeze(0) if attention_mask.ndim == 1 else attention_mask
        if input_ids.shape[-1] > frames:
            raise ValueError(
                f"The text stream holds {input_ids.shape[-1]} tokens for {frames} audio frames; the model reads one "
                "text token per frame."
            )
        if input_ids.shape[-1] < frames:
            padding = input_ids.new_full((input_ids.shape[0], frames - input_ids.shape[-1]), self.tokenizer.pad_token_id)
            input_ids = torch.cat([input_ids, padding], dim=-1)
            # Padding frames are predicted like any other frame, so they stay attended.
            attention_mask = None if attention_mask is None else torch.cat([attention_mask, torch.ones_like(padding)], -1)
        return input_ids, attention_mask

    def _persona_prefix(self, voice_prompt, text_prompt, merged) -> dict:
        config = self.audio_tokenizer.config
        frame_size = int(config.sampling_rate / config.frame_rate)
        silence_frames = int(SILENCE_DURATION * config.frame_rate)
        if voice_prompt is not None:
            waveform = self._normalize_voice(voice_prompt, config.sampling_rate, merged)
            voice_frames = int(np.ceil(waveform.shape[-1] / frame_size))
            waveform = np.pad(waveform, (0, voice_frames * frame_size - waveform.shape[-1]))
        else:
            voice_frames = 0
        prompt_ids = self._encode_text_prompt(text_prompt)
        total = voice_frames + silence_frames + prompt_ids.shape[-1] + silence_frames

        assistant = self._repeat_frame(self.silence_codes, total)
        if voice_frames:
            values = torch.as_tensor(waveform, dtype=torch.float32, device=self.audio_tokenizer.device)
            assistant[:, :, :voice_frames] = self._encode(values.reshape(1, 1, -1))
        input_ids = prompt_ids.new_full((1, total), self.tokenizer.pad_token_id)
        start = voice_frames + silence_frames
        input_ids[:, start : start + prompt_ids.shape[-1]] = prompt_ids
        return {
            "input_ids": input_ids,
            "attention_mask": torch.ones_like(input_ids),
            "user_audio_codes": self._repeat_frame(self.sine_codes, total),
            "assistant_audio_codes": assistant,
        }

    def _encode_text_prompt(self, text_prompt: str | None) -> torch.Tensor:
        if text_prompt is None:
            return torch.empty((1, 0), dtype=torch.long)
        cleaned = text_prompt.strip()
        if not (cleaned.startswith("<system>") and cleaned.endswith("<system>")):
            cleaned = f"<system> {cleaned} <system>"
        # `<system>` is ordinary text, cut into several tokens exactly as in training.
        return self.tokenizer(cleaned, return_tensors="pt", add_special_tokens=False)["input_ids"]

    def _normalize_voice(self, voice_prompt, sampling_rate: int, merged) -> np.ndarray:
        values = self.feature_extractor(voice_prompt, **merged["audio_kwargs"])["input_values"]
        waveform = np.asarray(values, dtype=np.float32).reshape(-1)
        try:
            import pyloudnorm
        except ImportError as error:
            raise ImportError("The voice prompt is loudness-normalized with `pyloudnorm`: `pip install pyloudnorm`.") from error
        loudness = pyloudnorm.Meter(sampling_rate).integrated_loudness(waveform)
        return np.asarray(pyloudnorm.normalize.loudness(waveform, loudness, VOICE_PROMPT_LUFS), dtype=np.float32)

    def _repeat_frame(self, codes: list[int], frames: int) -> torch.Tensor:
        frame = torch.tensor(codes, dtype=torch.long, device=self.audio_tokenizer.device)
        return frame.view(1, -1, 1).expand(1, -1, frames).contiguous()

    def _encode(self, values: torch.Tensor) -> torch.Tensor:
        values = values.to(self.audio_tokenizer.device, self.audio_tokenizer.dtype)
        return self.audio_tokenizer.encode(values, num_quantizers=self.num_codebooks).audio_codes

    def get_silence_audio_codes(self, num_frames: int, batch_size: int = 1) -> torch.Tensor:
        r"""
        Encodes `num_frames` frames of silence. Mimi is a streaming codec, so silence does not quantize to one
        repeated frame; the whole span is encoded at once.

        Args:
            num_frames (`int`):
                Frames to produce.
            batch_size (`int`, *optional*, defaults to 1):
                Identical rows to return.

        Returns:
            `torch.Tensor` of shape `(batch_size, num_codebooks, num_frames)`.
        """
        config = self.audio_tokenizer.config
        silence = torch.zeros((1, 1, num_frames * int(config.sampling_rate / config.frame_rate)))
        return self._encode(silence).expand(batch_size, -1, -1)

    def decode_audio(self, audio_codes: torch.Tensor, **kwargs) -> torch.Tensor:
        r"""
        Args:
            audio_codes (`torch.Tensor` of shape `(batch_size, num_codebooks, frames)`):
                Generated audio codes.

        Returns:
            `torch.Tensor` of shape `(batch_size, 1, samples)`: The waveform at 24 kHz.
        """
        return self.audio_tokenizer.decode(audio_codes.to(self.audio_tokenizer.device), **kwargs).audio_values


__all__ = ["PolyplexProcessor"]