import os
from modeling_videochat2 import *
from modeling_base import freeze_module
from transformers import AutoConfig
token = os.environ['HF_TOKEN']


class InternVideo2_cls(InternVideo2_VideoChat2):
    def __init__(self, config):
        super(InternVideo2_VideoChat2, self).__init__(config=config)

    def build_llm(self):
        self.lm_name = self.model_config.llm.name
        if self.model_config.llm.name == 'mistral_7b':
            from transformers import AutoModelForSequenceClassification
            config = AutoConfig.from_pretrained(
                self.model_config.llm.pretrained_llm_path,
                torch_dtype=torch.bfloat16,
                token=token,
                num_labels=self.model_config.llm.num_labels
                # attn_implementation="flash_attention_2",
            )
            self.lm = AutoModelForSequenceClassification.from_config(config)

        else:
            raise NotImplementedError(self.model_config.llm.name)

        self.freeze_llm = self.model_config.get("freeze_llm", True)
        logger.info(f'freeze_llm: {self.freeze_llm}')
        if self.freeze_llm:
            logger.info("freeze llm")
            freeze_module(self.lm)
        
        if self.model_config.llm.use_lora:
            self.use_lora = True
            from peft import get_peft_model, LoraConfig, TaskType
            logger.info("Use lora")

            peft_config = LoraConfig(
                task_type=TaskType.CAUSAL_LM, inference_mode=False, 
                r=self.model_config.llm.lora_r, lora_alpha=self.model_config.llm.lora_alpha, lora_dropout=self.model_config.llm.lora_dropout,
                target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                                "gate_proj", "up_proj", "down_proj"]
                )
                
            self.lm = get_peft_model(self.lm, peft_config)
            self.lm.enable_input_require_grads()
            self.lm.print_trainable_parameters()
        else:
            self.use_lora = False


    def build_conversation(self,instruction, user_prompt,media_type='video',msg=''):

        conversation = ""
        if instruction:
            conversation += instruction
        conversation += ("[INST]" + " ")

        if media_type == 'image':
            conversation +=( "<Image>" + IMG_TOKEN + "</Image>")#*ilen
        else:
            conversation += ("<Video>" + VID_TOKEN + "</Video>")#*ilen

        conversation += (msg.rstrip() + "[/INST]")
        conversation += (" [INST] " + user_prompt + " [/INST]")
        conversation += ("")
        return conversation

    def test(self, x):
        return x

if __name__ == "__main__":

    tokenizer =  AutoTokenizer.from_pretrained('OpenGVLab/InternVideo2-Chat-8B',trust_remote_code=True,use_fast=False)
    config = AutoConfig.from_pretrained('OpenGVLab/InternVideo2-Chat-8B', torch_dtype=torch.bfloat16,trust_remote_code=True)
    model = InternVideo2_Classification(config).cuda()

    B, T, C, H, W = 1, 8, 3, 224, 224
    video_tensor = torch.randn(B,T,C,H,W).cuda()
    user_prompt = "this is a user prompt"
    instruction = "this is an instruction"

    conversation = model.build_conversation(instruction=instruction, user_prompt=user_prompt, media_type='video')
    tokenized = model.build_input_ids(tokenizer,conversation,max_length=248,add_special_tokens=True,truncation=False,padding=False,return_tensors='pt')
    
    input_ids = tokenized['input_ids'].unsqueeze(0).to(model.device)
    attn_mask = tokenized['attention_mask'].unsqueeze(0).to(model.device)
    indexes = tokenized['index'].unsqueeze(0)
    text_embeds = model.pad_text_embeds(input_ids = input_ids,video = video_tensor,video_idx = indexes)
    outputs = model.lm(inputs_embeds=text_embeds, attention_mask=attn_mask,output_hidden_states=True,return_dict=True)