from __future__ import annotations

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

from .configuration_tinyaya import TinyAyaConfig


class TinyAyaForCausalLM(Cohere2ForCausalLM):
    config_class = TinyAyaConfig

    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),
        )

    @staticmethod
    def _sample(scores, do_sample, temperature, top_k):
        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,
        max_new_tokens=2048,
        min_new_tokens=8,
        temperature=0.8,
        top_k=30,
        do_sample=True,
    ):
        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._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