from llava.model.builder import load_pretrained_model
from llava.utils import disable_torch_init
from llava.mm_utils import (
    tokenizer_image_token,
    process_images,
    get_model_name_from_path,
)
from llava.constants import DEFAULT_IMAGE_TOKEN, DEFAULT_AUDIO_TOKEN
from llava.mm_utils import tokenizer_image_token
from llava.action_tokenizer import ActionTokenizer, encode_robot_obs ,denormalize_actions
from llava import conversation as conversation_lib

import torch
import numpy as np

class LLaVA:
    
    def __init__(self, args):
        # Initialize model with provided arguments
        model_path = args.model_path
        model_name = get_model_name_from_path(model_path)
        model_base = args.model_base
        self.tokenizer, self.llm_robot, self.image_processor, self.context_len = (
            load_pretrained_model(model_path, model_base, model_name)
        )
        self.temperature = args.temperature
        self.top_p = args.top_p
        self.num_beams = args.num_beams
        self.max_new_tokens = args.max_new_tokens
        self.action_tokenizer = ActionTokenizer(self.tokenizer)
        self.action_stat = args.action_stat

    # set language randomly
    def set_language(self, instruction):
        self.instruction = instruction

    def compose_robot_input(self, image, instruction, robot_obs):
        # Prepare input data for the robot model
        image = self.image_processor.preprocess(image, return_tensors="pt")["pixel_values"][0]
        image = image[None, :]
        image = image.to(dtype=torch.float16, device="cuda", non_blocking=True)
  
        robot_obs = [str(elem) for elem in robot_obs]
        robot_obs = " ".join(robot_obs)
        robot_obs = encode_robot_obs(robot_obs, self.action_tokenizer, self.action_stat)
        
        instruction = DEFAULT_IMAGE_TOKEN + "\n" + instruction + "\n" + robot_obs
        conv = conversation_lib.default_conversation.copy()
        conv.system = "A chat between a curious user and an artificial intelligence robot. The robot provides actions to follow out the user's instructions."
        conv.append_message(conv.roles[0], instruction)
        conv.append_message(conv.roles[1], None)
        instruction = conv.get_prompt()
        
        input_ids = tokenizer_image_token(instruction, self.tokenizer, return_tensors="pt")
        input_ids = torch.stack([input_ids],dim=0)
        
        return input_ids, image
    
    def get_action(self, input_ids, images):
       # Generate actions based on input IDs and images
        with torch.inference_mode():
            output_ids = self.llm_robot.generate(
                input_ids.cuda(),
                images=images.to(dtype=torch.float16, device="cuda", non_blocking=True),
                do_sample=True if self.temperature > 0 else False,
                temperature=self.temperature,
                top_p=self.top_p,
                num_beams=self.num_beams,
                max_new_tokens=self.max_new_tokens,
                use_cache=True,
                )            
            
            output_ids = output_ids[0].cpu().numpy().tolist()[2:-1]
            actions = [self.action_tokenizer.decode_token_ids_to_actions(elem) for elem in output_ids]
            actions = denormalize_actions(actions,self.action_stat)

            return np.array(actions)