# Copyright (c) ModelScope Contributors. All rights reserved.
import itertools
import torch
from functools import partial
from typing import Any, Dict, List, Literal, Optional

from swift.utils import get_env_args, to_float_dtype
from ..base import Template
from ..constant import MLLMTemplateType
from ..register import TemplateMeta, register_template
from ..template_inputs import StdTemplateInputs
from ..utils import Context, findall
from ..vision_utils import load_batch, load_file
from .qwen import QwenTemplateMeta


class GOTImageEvalProcessor:

    def __init__(self, image_size=384, mean=None, std=None):
        from torchvision import transforms
        from torchvision.transforms.functional import InterpolationMode
        if mean is None:
            mean = (0.48145466, 0.4578275, 0.40821073)
        if std is None:
            std = (0.26862954, 0.26130258, 0.27577711)

        self.normalize = transforms.Normalize(mean, std)

        self.transform = transforms.Compose([
            transforms.Resize((image_size, image_size), interpolation=InterpolationMode.BICUBIC),
            transforms.ToTensor(),
            self.normalize,
        ])

    def __call__(self, item):
        return self.transform(item)


class GOT_OCR2Template(Template):
    placeholder_tokens = ['<imgpad>']

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        # 'OCR: '
        # 'OCR with format: '
        assert media_type == 'image'
        return ['<img>' + '<imgpad>' * 256 + '</img>\n']

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = super()._encode(inputs)
        images = inputs.images
        image_processor_high = GOTImageEvalProcessor(image_size=1024)
        for i, image in enumerate(images):
            images[i] = image_processor_high(image)[None].to(self.model_info.torch_dtype)
        if images:
            encoded['images'] = images
        return encoded

    def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
        res = super()._data_collator(batch, padding_to=padding_to)
        images = self.gather_list(batch, 'images')
        if images:
            res['images'] = images
        return res


register_template(
    QwenTemplateMeta(
        MLLMTemplateType.got_ocr2,
        default_system='        You should follow the instructions carefully and explain your answers in detail.',
        template_cls=GOT_OCR2Template,
        agent_template=None,
    ))


class GOT_OCR2HfTemplate(Template):
    placeholder_tokens = ['<imgpad>']

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        # 'OCR: '
        # 'OCR with format: '
        assert media_type == 'image'
        return ['<img>' + '<imgpad>' * 256 + '</img>\n']

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:  # 暂时照抄上面
        encoded = super()._encode(inputs)
        images = inputs.images
        if images:
            encoded['images'] = images
        return encoded

    def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
        res = super()._data_collator(batch, padding_to=padding_to)
        images = self.gather_list(batch, 'images')
        _inputs = self.processor(images, return_tensors='pt')
        _inputs.pop('input_ids')  # this does not contain the response, so cannot be used when training
        _inputs.pop('attention_mask')  # this does not contain the response, so cannot be used when training

        res.update(_inputs.data)
        return res


register_template(
    QwenTemplateMeta(
        MLLMTemplateType.got_ocr2_hf,
        default_system='        You should follow the instructions carefully and explain your answers in detail.',
        template_cls=GOT_OCR2HfTemplate,
        agent_template=None,
    ))


class StepAudioTemplate(Template):
    use_model = True

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        assert media_type == 'audio', f'media_type: {media_type}'
        from utils import load_audio
        audio_wav, sr = load_audio(load_file(inputs.audios[index]))
        audio_tokens = self.model.encoder(audio_wav, sr)
        return audio_tokens


class StepAudio2MiniTemplate(Template):
    use_model = True

    def load_audio(self, file_path, target_rate=16000, max_length=None):
        '''
        Open an audio file and read as mono waveform, resampling as necessary
        If max_length is provided, truncate the audio to that length
        '''
        import torchaudio
        waveform, sample_rate = torchaudio.load(file_path)
        if sample_rate != target_rate:
            waveform = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=target_rate)(waveform)
        audio = waveform[0]  # get the first channel

        # Truncate audio if it exceeds max_length
        if max_length is not None and audio.shape[0] > max_length:
            audio = audio[:max_length]

        return audio

    def _mel_filters(self, n_mels: int) -> 'torch.Tensor':
        '''Load the mel filterbank matrix for projecting STFT into a Mel spectrogram.'''
        import librosa
        import torch
        assert n_mels in {80, 128}, f'Unsupported n_mels: {n_mels}'
        if n_mels == 128:
            return torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=128))
        else:
            return torch.from_numpy(librosa.filters.mel(sr=16000, n_fft=400, n_mels=80))

    def log_mel_spectrogram(self, audio, n_mels=128, padding=479):
        '''
        Compute the log-Mel spectrogram with specific padding for StepAudio
        '''
        import torch
        import torch.nn.functional as F
        if isinstance(audio, str):
            audio = self.load_audio(audio)
        elif not torch.is_tensor(audio):
            audio = torch.from_numpy(audio)
        if padding > 0:
            audio = F.pad(audio, (0, padding))
        window = torch.hann_window(400).to(audio.device)
        stft = torch.stft(audio, 400, 160, window=window, return_complex=True)
        magnitudes = stft[..., :-1].abs()**2
        filters = self._mel_filters(n_mels)
        mel_spec = filters @ magnitudes

        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
        log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
        log_spec = (log_spec + 4.0) / 4.0
        return log_spec

    def compute_token_num(self, max_feature_len):
        # First, audio goes through encoder:
        # 1. conv1: kernel=3, stride=1, padding=1 -> size unchanged
        # 2. conv2: kernel=3, stride=2, padding=1 -> size/2
        # 3. avg_pooler: kernel=2, stride=2 -> size/2
        max_feature_len = max_feature_len - 2  # remove padding
        encoder_output_dim = (max_feature_len + 1) // 2 // 2  # after conv2 and avg_pooler

        # Then through adaptor (parameters from config file):
        padding = 1
        kernel_size = 3  # from config: audio_encoder_config.kernel_size
        stride = 2  # from config: audio_encoder_config.adapter_stride
        adapter_output_dim = (encoder_output_dim + 2 * padding - kernel_size) // stride + 1
        return adapter_output_dim

    def padding_mels(self, data: List['torch.Tensor']):
        ''' Padding the data into batch data

        Parameters
        ----------
            data: List[Tensor], shape of Tensor (128, T)

        Returns:
        -------
            feats, feats lengths
        '''
        import torch
        from torch.nn.utils.rnn import pad_sequence
        sample = data
        assert isinstance(sample, list)
        feats_lengths = torch.tensor([s.size(1) - 2 for s in sample], dtype=torch.int32)
        feats = [s.t() for s in sample]
        padded_feats = pad_sequence(feats, batch_first=True, padding_value=0)

        return padded_feats.transpose(1, 2), feats_lengths

    def audio_process(self, audio):
        results = []
        mels = []
        for i in range(0, audio.shape[0], 16000 * 25):
            mel = self.log_mel_spectrogram(audio[i:i + 16000 * 25], n_mels=128, padding=479)
            mels.append(mel)
            audio_tokens = '<audio_patch>' * self.compute_token_num(mel.shape[1])
            results.append(f'<audio_start>{audio_tokens}<audio_end>')
        audio_ids = self._tokenize(''.join(results))
        return audio_ids, mels

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        assert media_type == 'audio'
        return ['<audio_patch>']

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = super()._encode(inputs)
        input_ids = encoded['input_ids']
        labels = encoded['labels']
        loss_scale = encoded.get('loss_scale', None)
        sampling_rate = get_env_args('sampling_rate', int, 16000)
        inputs.audios = load_batch(inputs.audios, partial(self.load_audio, target_rate=sampling_rate))

        audio_token = self._tokenize('<audio_patch>')[0]
        idx_list = findall(input_ids, audio_token)

        if idx_list:
            audio_inputs = []
            mels = []
            for audio in inputs.audios:
                audio_input, mel = self.audio_process(audio)
                audio_inputs.append(audio_input)
                mels.extend(mel)

            def _get_new_audio_tokens(i):
                return audio_inputs[i]

            input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
                                                                _get_new_audio_tokens)
            encoded['input_ids'] = input_ids  # Add labels to the batch
            encoded['labels'] = labels  # Add labels to the batch
            encoded['loss_scale'] = loss_scale
            encoded['mels'] = mels
            wavs, wav_lens = self.padding_mels(mels)
            # audio_tokens = [151688, 151690, 151689]
            # for audio_token_id in audio_tokens:
            #     labels[labels == audio_token_id] = -100  # Mask image token IDs in labels

        else:
            wavs = None
            wav_lens = None

        encoded['wavs'] = wavs
        encoded['wav_lens'] = wav_lens

        return encoded

    def _data_collator(self, batch: List[Dict[str, Any]], *, padding_to: Optional[int] = None) -> Dict[str, Any]:
        combined_mels = list(itertools.chain.from_iterable([e.get('mels', []) for e in batch]))
        batch_wavs, batch_wav_lens = self.padding_mels(combined_mels) if combined_mels else (None, None)
        res = super()._data_collator(batch, padding_to=padding_to)
        res['wav_lens'] = batch_wav_lens
        res['wavs'] = batch_wavs

        return res


register_template(
    TemplateMeta(
        MLLMTemplateType.step_audio2_mini,
        template_cls=StepAudio2MiniTemplate,
        prefix=[],
        prompt=['<|BOT|>human\n{{QUERY}}<|EOT|><|BOT|>assistant\n'],
        system_prefix=['<|BOT|>system\n{{SYSTEM}}<|EOT|>'],
        chat_sep=['<|EOT|>'],
        suffix=['<|EOT|>'],
    ))

register_template(
    TemplateMeta(
        MLLMTemplateType.step_audio,
        template_cls=StepAudioTemplate,
        prefix=['<s>'],
        prompt=['<|BOT|>human\n{{QUERY}}<|EOT|><|BOT|>assistant\n'],
        system_prefix=['<s><|BOT|>system\n{{SYSTEM}}<|EOT|>'],
        chat_sep=['<|EOT|>'],
        suffix=['<|EOT|>'],
    ))


class Step3VLTemplate(Template):
    use_model = True
    support_padding_free = False
    image_token_id = 151679
    placeholder_tokens = ['<im_patch>']

    def replace_tag(self, media_type: Literal['image', 'video', 'audio'], index: int,
                    inputs: StdTemplateInputs) -> List[Context]:
        assert media_type == 'image'
        return ['<im_patch>']

    def _encode(self, inputs: StdTemplateInputs) -> Dict[str, Any]:
        encoded = super()._encode(inputs)
        input_ids = encoded['input_ids']
        labels = encoded['labels']
        loss_scale = encoded.get('loss_scale', None)

        images = inputs.images
        if images:
            processor = self.processor
            idx_list = findall(input_ids, self.image_token_id)
            splitted_images_data = processor._split_images(images)
            pixel_values_lst = []
            patch_pixel_values_lst = []
            patch_newline_mask_lst = []
            image_repl_ids_lst = []
            num_patches = []
            for raw_img, img_patches, patch_newline_mask in splitted_images_data:
                pixel_values_lst.extend(processor._convert_images_to_pixel_values([raw_img]))
                if len(img_patches) > 0:
                    patch_pixel_values_lst.extend(processor._convert_images_to_pixel_values(img_patches, is_patch=True))
                num_patches.append(len(img_patches))
                _, image_repl_ids = processor._get_image_repl_features(1, len(img_patches), patch_newline_mask)
                image_repl_ids_lst.append(image_repl_ids)
                if patch_newline_mask is not None:
                    patch_newline_mask_lst.extend(patch_newline_mask)

            image_inputs = {
                'pixel_values': torch.cat(pixel_values_lst),
                'num_patches': num_patches,
            }
            if patch_pixel_values_lst:
                image_inputs['patch_pixel_values'] = torch.cat(patch_pixel_values_lst)
            if patch_newline_mask_lst:
                image_inputs['patch_newline_mask'] = torch.tensor(patch_newline_mask_lst, dtype=torch.bool)
            image_inputs = to_float_dtype(image_inputs, self.model_info.torch_dtype)

            def _get_new_tokens(i):
                return image_repl_ids_lst[i]

            input_ids, labels, loss_scale = self._extend_tokens(input_ids, labels, loss_scale, idx_list,
                                                                _get_new_tokens)

        encoded['input_ids'] = input_ids
        encoded['labels'] = labels
        encoded['loss_scale'] = loss_scale
        encoded.update(image_inputs)
        return encoded

    def _post_encode(self, model, inputs: Dict[str, Any]) -> Dict[str, Any]:
        if not self.is_training:
            return inputs

        input_ids = inputs['input_ids']
        # Only one image is supported per sample. # File: modeling_step_vl.py line 319, in _process_image_input
        #           cur_feature.append(image_features[i].view(-1, image_features.shape[-1]))
        pixel_values = inputs.get('pixel_values')
        num_patches = inputs.get('num_patches')
        patch_pixel_values = inputs.get('patch_pixel_values')

        base_model = self.get_base_model(model)
        inputs_embeds = base_model.model.language_model.embed_tokens(input_ids)
        if pixel_values is not None:
            img_inputs = base_model.model._parse_and_validate_image_input(
                pixel_values=pixel_values, num_patches=num_patches, patch_pixel_values=patch_pixel_values)

            # [image embedding or concatenation of image embedding and patch image embedding]
            img_embeddings = base_model.model._process_image_input(img_inputs)

            is_multimodal = input_ids == self.image_token_id
            is_multimodal = is_multimodal.to(inputs_embeds.device)
            bs = is_multimodal.shape[0]
            for i in range(bs):
                assert is_multimodal[i].sum() == img_embeddings[i].shape[0]
            B, L, D = inputs_embeds.shape
            flat_img_embeds = torch.cat(img_embeddings, dim=0)
            flat_mask = is_multimodal.view(-1)
            flat_inputs_embeds = inputs_embeds.view(-1, D)
            flat_inputs_embeds[flat_mask] = flat_img_embeds
            inputs_embeds = flat_inputs_embeds.view(B, L, D)
        return {'inputs_embeds': inputs_embeds}

    def _data_collator_mm_data(self, batch: List[Dict[str, Any]]) -> Dict[str, Any]:
        res = super()._data_collator_mm_data(batch)
        num_patches = self.gather_list(batch, 'num_patches')
        if num_patches:
            res['num_patches'] = num_patches
        patch_pixel_values = [b['patch_pixel_values'] for b in batch if b.get('patch_pixel_values') is not None]
        patch_newline_mask = [b['patch_newline_mask'] for b in batch if b.get('patch_newline_mask') is not None]
        if patch_pixel_values:
            res['patch_pixel_values'] = torch.concat(patch_pixel_values)
            res['patch_newline_mask'] = torch.concat(patch_newline_mask)
        return res


register_template(
    QwenTemplateMeta(
        MLLMTemplateType.step3_vl,
        template_cls=Step3VLTemplate,
        default_system=None,
        is_thinking=True,
        thinking_prefix='<think>\n',
    ))