"""Configuration for rumik-oss 1."""

from transformers import Cohere2Config


class RumikOSSConfig(Cohere2Config):
    """rumik-oss 1: tiny aya fire extended with flattened Mimi codec tokens.

    Beyond the base Cohere2 fields, this records where the audio vocabulary
    lives so that callers can map token ids to codec (code, quantizer) pairs
    arithmetically, without loading the tokenizer or scanning its vocabulary.

    The unit tokens are laid out code-major, quantizer-minor::

        <0_0> <0_1> ... <0_7> <1_0> ... <2047_7>

    so for any id in ``[first_unit_id, last_unit_id]``::

        code      = (token_id - first_unit_id) // num_quantizers
        quantizer = (token_id - first_unit_id) %  num_quantizers
    """

    model_type = "rumik_oss"

    def __init__(
        self,
        num_quantizers=8,
        codebook_size=2048,
        audio_start_token_id=None,
        audio_end_token_id=None,
        text_start_token_id=None,
        first_unit_id=None,
        last_unit_id=None,
        frame_rate_hz=12.5,
        speakers=None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.num_quantizers = int(num_quantizers)
        self.codebook_size = int(codebook_size)
        self.audio_start_token_id = audio_start_token_id
        self.audio_end_token_id = audio_end_token_id
        self.text_start_token_id = text_start_token_id
        self.first_unit_id = first_unit_id
        # Derivable, but stored so a hand-edited config stays self-consistent.
        if last_unit_id is None and first_unit_id is not None:
            last_unit_id = int(first_unit_id) + self.codebook_size * self.num_quantizers - 1
        self.last_unit_id = last_unit_id
        self.frame_rate_hz = float(frame_rate_hz)
        self.speakers = speakers or ["Ira", "Aisha", "Siya", "Zoya"]

    @property
    def audio_tokens_per_second(self) -> float:
        """Tokens the model emits per second of audio (8 x 12.5 = 100)."""
        return self.frame_rate_hz * self.num_quantizers