"""rumik-oss 1: text -> flattened Mimi codec tokens.

This model produces *audio tokens*, not waveforms. Decoding them to audio is
the caller's job and needs the Mimi codec (shipped in ``codec/``)::

    ids    = model.generate_audio(**inputs)                  # prompt + audio
    tokens = ids[0].tolist()[inputs.input_ids.shape[1]:]     # drop the prompt
    codes  = model.audio_tokens_to_codes(tokens)             # [1, 8, T] for Mimi
    wav    = mimi.decode(codes).audio_values                 # your call
"""
from __future__ import annotations

import torch
from torch import nn
from transformers.models.cohere2.modeling_cohere2 import Cohere2ForCausalLM

from .configuration_rumik_oss import RumikOSSConfig


class RumikOSSForCausalLM(Cohere2ForCausalLM):
    config_class = RumikOSSConfig

    def __init__(self, config):
        super().__init__(config)
        hidden = int(config.hidden_size)
        self.stop_predictor = nn.Sequential(
            nn.LayerNorm(hidden),
            nn.Linear(hidden, max(64, hidden // 4)),
            nn.GELU(),
            nn.Linear(max(64, hidden // 4), 1),
        )

    # ---- audio vocabulary -------------------------------------------------
    def audio_token_ids(self, device=None) -> torch.Tensor:
        """Every id the model may legally emit inside an <audio> span.

        Derived from the config, so callers no longer have to scan the
        tokenizer vocabulary for ``<code>_<quantizer>`` tokens.
        """
        c = self.config
        if c.first_unit_id is None:
            raise ValueError(
                "config.first_unit_id is unset; pass allowed_ids explicitly or "
                "add first_unit_id/last_unit_id to config.json.")
        ids = torch.arange(int(c.first_unit_id), int(c.last_unit_id) + 1, dtype=torch.long)
        end = torch.tensor([int(c.audio_end_token_id)], dtype=torch.long)
        return torch.cat((ids, end)).to(device if device is not None else self.device)

    def audio_tokens_to_codes(self, token_ids) -> torch.Tensor:
        """Generated ids -> ``[1, num_quantizers, num_frames]`` codec codes.

        Shaped for ``MimiModel.decode``. Purely arithmetic -- no tokenizer, no
        codec, no audio. Pass only the generated tail, not the prompt.

        A token that breaks the round robin does not end the audio: the
        partial frame is dropped and the next complete frame is picked up
        again, so one bad token costs one frame rather than the rest of the
        clip. ``</audio>`` ends it.
        """
        c = self.config
        first, last, Q = int(c.first_unit_id), int(c.last_unit_id), int(c.num_quantizers)
        end_id = int(c.audio_end_token_id)
        if torch.is_tensor(token_ids):
            token_ids = token_ids.flatten().tolist()

        frames: list[list[int]] = []
        frame: list[int] = []
        for tid in (int(t) for t in token_ids):
            if tid == end_id:
                break
            if not first <= tid <= last:
                frame = []                    # stray token: resync
                continue
            code, q = divmod(tid - first, Q)
            if q == len(frame):
                frame.append(code)
                if len(frame) == Q:
                    frames.append(frame)
                    frame = []
            else:                             # off the round robin
                frame = [code] if q == 0 else []
        if not frames:
            raise ValueError(
                f"no complete codec frame in {len(token_ids)} tokens "
                f"(need at least {Q})")
        return torch.tensor(frames, dtype=torch.long).T.unsqueeze(0)

    # ---- generation -------------------------------------------------------
    @staticmethod
    def _constrained_sample(scores, do_sample, temperature, top_k):
        # NB: deliberately not named `_sample`. `GenerationMixin._sample` is the
        # method `generate()` dispatches to, and shadowing it makes every call
        # to `generate()` fail with a TypeError on `logits_processor`.
        if not do_sample:
            return scores.argmax(dim=-1, keepdim=True)
        scores = scores / max(float(temperature), 1e-5)
        if int(top_k) > 0:
            k = min(int(top_k), scores.shape[-1])
            cutoff = torch.topk(scores, k, dim=-1).values[:, -1:]
            scores = scores.masked_fill(scores < cutoff, torch.finfo(scores.dtype).min)
        return torch.multinomial(torch.softmax(scores.float(), dim=-1), 1)

    @torch.inference_mode()
    def generate_audio(
        self,
        input_ids,
        attention_mask,
        allowed_ids=None,
        max_new_tokens=2048,
        min_new_tokens=8,
        temperature=0.8,
        top_k=30,
        do_sample=True,
    ):
        """Autoregressively emit audio tokens, ending at ``</audio>``.

        ``allowed_ids`` defaults to the audio vocabulary from the config.
        """
        if allowed_ids is None:
            allowed_ids = self.audio_token_ids(device=input_ids.device)
        out = self(
            input_ids=input_ids,
            attention_mask=attention_mask,
            use_cache=True,
            output_hidden_states=True,
            return_dict=True,
        )
        emitted = input_ids
        cache = out.past_key_values
        scores = out.logits[:, -1, :]
        hidden = out.hidden_states[-1][:, -1:, :]
        mask = attention_mask
        allowed_ids = allowed_ids.to(scores.device)
        audio_end_id = int(self.config.audio_end_token_id)
        for step in range(int(max_new_tokens)):
            if step >= int(min_new_tokens):
                stop = torch.sigmoid(self.stop_predictor(hidden).squeeze(-1))
                if bool((stop > 0.5).all()):
                    eos = input_ids.new_full((input_ids.shape[0], 1), audio_end_id)
                    return torch.cat((emitted, eos), dim=1)
            restricted = torch.full_like(scores, torch.finfo(scores.dtype).min)
            restricted.index_copy_(1, allowed_ids, scores.index_select(1, allowed_ids))
            if step < int(min_new_tokens):
                restricted[:, audio_end_id] = torch.finfo(scores.dtype).min
            token = self._constrained_sample(restricted, do_sample, temperature, top_k)
            emitted = torch.cat((emitted, token), dim=1)
            if bool((token == audio_end_id).all()):
                break
            mask = torch.cat((mask, torch.ones_like(token)), dim=1)
            out = self(
                input_ids=token,
                attention_mask=mask,
                past_key_values=cache,
                use_cache=True,
                output_hidden_states=True,
                return_dict=True,
            )
            cache = out.past_key_values
            scores = out.logits[:, -1, :]
            hidden = out.hidden_states[-1][:, -1:, :]
        return emitted