import importlib

import mmengine
import torch
import torch.nn as nn
from mmengine.device import get_device

from opencompass.registry import MM_MODELS


@MM_MODELS.register_module('otter-9b')
class Otter(nn.Module):
    """Inference code of OTTER.

    Model details:
        OTTER: a multi-modal model based on OpenFlamingo
        (open-sourced version of DeepMind's Flamingo)
        https://github.com/Luodian/Otter
    Args:
        model_path (str): The path of OTTER model
        in Huggingface model hub format.
        load_bit (str): The bit of OTTER model, can be "fp32" or "bf16".
        mode (str): The mode of inference. Defaults to 'generation'.
    """

    def __init__(self,
                 model_path,
                 load_bit,
                 prompt_constructor,
                 post_processor,
                 mode='generation') -> None:
        super().__init__()
        torch_dtype = torch.bfloat16 if load_bit == 'bf16' else torch.float32
        otter_ai = importlib.import_module('otter_ai')
        self.model = otter_ai.OtterForConditionalGeneration.from_pretrained(
            model_path, torch_dtype=torch_dtype, device_map=get_device())
        self.tokenizer = self.model.text_tokenizer
        self.tokenizer.padding_side = 'left'
        self.model_dtype = next(self.model.parameters()).dtype
        self.prompt_constructor = mmengine.registry.build_from_cfg(
            prompt_constructor, MM_MODELS)
        if post_processor is not None:
            self.post_processor = mmengine.registry.build_from_cfg(
                post_processor, MM_MODELS)
        self.mode = mode

    def forward(self, batch):
        if self.mode == 'generation':
            return self.generate(batch)
        elif self.mode == 'loss':
            return self.loss(batch)
        else:
            raise RuntimeError(f'Invalid mode "{self.mode}".')

    def generate(self, batch):
        inputs = self.prompt_constructor(batch)
        image = inputs['image']
        prompt = inputs['prompt']
        data_samples = inputs['data_samples']
        vision_x = image.unsqueeze(1).unsqueeze(0).to(dtype=self.model_dtype)
        lang_x = self.model.text_tokenizer([prompt], return_tensors='pt')
        bad_words_id = self.model.text_tokenizer(['User:', 'GPT:']).input_ids
        generated_text = self.model.generate(
            vision_x=vision_x.to(self.model.device),
            lang_x=lang_x['input_ids'].to(self.model.device),
            attention_mask=lang_x['attention_mask'].to(self.model.device),
            do_sample=False,
            max_new_tokens=512,
            num_beams=3,
            bad_words_ids=bad_words_id,
            no_repeat_ngram_size=3,
        )
        for i, data_sample in enumerate(data_samples):
            output_text = self.post_processor(generated_text[i],
                                              self.model.text_tokenizer)
            data_sample.pred_answer = output_text
            data_samples[i] = data_sample

        return data_samples