import itertools

from typing import List, Optional, Tuple, Union
import torch
from torch import nn
from transformers.modeling_outputs import BaseModelOutput
from transformers import Wav2Vec2BertModel, MllamaPreTrainedModel
from .configuration_llama3 import MllamaAudioConfig


class MllamaAudioModel(MllamaPreTrainedModel):
    config_class = MllamaAudioConfig
    base_model_prefix = "audio_model"
    def __init__(self, config: MllamaAudioConfig, text_embedding: nn.Embedding):
        super().__init__(config)
        assert config.add_adapter is True, f'{type(self).__name__} requires add adapter to be true.'
        assert config.output_hidden_size == text_embedding.weight.shape[1], f'Output hidden size({config.output_hidden_size}) of audio model and text embedding({text_embedding.weight.shape[1]}) must match!'
        self.text_embedding = text_embedding
        self.audio_embedding = Wav2Vec2BertModel(config)
        self.start_of_audio = nn.Parameter(data=torch.mean(text_embedding.weight, dim=0).unsqueeze(0), requires_grad=True)
        self.end_of_audio = nn.Parameter(data=torch.mean(text_embedding.weight, dim=0).unsqueeze(0), requires_grad=True)
        self.filler_token_id = config.filler_token_id

    def forward(
        self,
        audio_features: torch.Tensor = None,
        input_ids: torch.LongTensor = None,
        return_dict: Optional[bool] = None,
    ) -> Union[BaseModelOutput, Tuple[torch.Tensor, ...]]:
        input_embeddings = self.text_embedding(torch.clamp(input_ids, min=0))
        bs, max_num_img, _, _ = audio_features.shape

        for i in range(bs):
            for j in range(max_num_img):
                audio_id = -1 - j
                idx = torch.where(input_ids[i] == audio_id)
                if idx.numel() > 0:
                    input_embeddings[i][idx] = torch.concat([self.start_of_audio, audio_features[i, j][idx], self.end_of_audio])

        idx = torch.where(input_ids < 0 and input_ids >= -max_num_img)
        input_ids[idx].fill_(self.filler_token_id)

        if return_dict:
            return dict(input_embeddings=input_embeddings)
        return input_embeddings