import math
from typing import Dict, List, Optional, Union
import numpy as np
import transformers
from transformers.tokenization_utils_base import AudioInput
from transformers.models.seamless_m4t.feature_extraction_seamless_m4t import SeamlessM4TFeatureExtractor
from transformers.utils import TensorType
from transformers.feature_extraction_utils import BatchFeature
from transformers import AutoFeatureExtractor


def build_audio_tokens(text: List[str], audio_features: Union[Dict, List[List[np.ndarray]]], audio_token="<|audio|>") -> Dict:
    if not isinstance(audio_features, list):
        audio_features = audio_features['audio_features']
    bs = audio_features.shape[0]
    for i in range(bs):
        for j in range(len(audio_features[i])):
            tgt_token = f"<|audio_{j+1}|>" * get_num_embeddings(audio_features[i][j].shape[0])
            text[i] = text[i].replace(audio_token, tgt_token, 1)
    return text
            
def get_num_embeddings(num_framses, adapter_kernel_size=7, adapter_stride=4) -> int:
    return math.ceil((num_framses - adapter_kernel_size) / adapter_stride) + 1 + 2 # 2 = <|begin_of_audio|>, <|end_of_audio|>


class MllamaAudioFeatureExtractor(SeamlessM4TFeatureExtractor):
    
    def __call__(
        self,
        batch_audio_clips: List[List[AudioInput]],
        return_tensors: Optional[Union[str, TensorType]] = None,
    ) -> BatchFeature:
        audio_features = [[ super(MllamaAudioFeatureExtractor, self).__call__(audio_j, sampling_rate=16000, return_attention_mask=False)['input_features'][0] for audio_j in audio_i ] for audio_i in batch_audio_clips ]
        packed_audio_features = self.pack_audio_clips(audio_features)
        
        encoded_audio_inputs = BatchFeature(
            data={
                "audio_features": packed_audio_features,
            },
            tensor_type=return_tensors,
        )
        
        return encoded_audio_inputs
    
    def pack_audio_clips(self, batch_audio_clips: List[List[np.ndarray]]) -> np.ndarray:
        assert batch_audio_clips[0][0].ndim == 2 # sequence length x feature dimension
        # Determine output shape: (batch_size, max_num_clips, max_frames, feature_dim)
        batch_size = len(batch_audio_clips)
        max_num_clips = max([len(clips) for clips in batch_audio_clips])
        max_frames = max([clip.shape[0] for clips in batch_audio_clips for clip in clips])
        feature_dim = batch_audio_clips[0][0].shape[1]
        
        stacked_audio_clips = np.zeros((batch_size, max_num_clips, max_frames, feature_dim), dtype=np.float32) 
        for i, clips in enumerate(batch_audio_clips):
            for j, clip in enumerate(clips):
                stacked_audio_clips[i, j, :clip.shape[0], :] = clip
            
        return stacked_audio_clips

AutoFeatureExtractor.register("MllamaAudioFeatureExtractor", MllamaAudioFeatureExtractor)
transformers.MllamaAudioFeatureExtractor = MllamaAudioFeatureExtractor