import torch
import cv2
from PIL import Image
from transformers import AutoModelForMaskedLM, AutoTokenizer, AutoModel, AutoConfig, AutoModelForMaskedLM
import numpy as np
CAMERA_VIEWS=['cam_bottom', 'cam_top', 'cam_left_wrist', 'cam_right_wrist']

from dex_vla.model_load_utils import load_model_for_eval
class paligemma_vla_policy:
    def __init__(self, policy_config, data_args=None):
        super(paligemma_vla_policy).__init__()
        self.load_policy(policy_config)
        self.history_len = policy_config['history_image_length']
        self.data_args = data_args

    def load_policy(self, policy_config):
        self.policy_config = policy_config
        # self.conv = conv_templates[policy_config['conv_mode']].copy()
        model_base = policy_config["model_base"] if policy_config[
            'enable_lora'] else None
        model_path = policy_config["model_path"]

        self.tokenizer, self.policy, self.multimodal_processor, self.context_len = load_model_for_eval(model_path=model_path,
                                                                                                    model_base=model_base, policy_config=policy_config)
        # self.tokenizer.add_special_tokens({'additional_special_tokens': ["[SOA]"]})

        self.config = AutoConfig.from_pretrained('/'.join(model_path.split('/')[:-1]), trust_remote_code=True)

    def process_batch_to_qwen2_vla(self, curr_image, robo_state, raw_lang):
        curr_image = curr_image[-self.history_len:]
        if len(curr_image) == 1 and self.history_len > 1:
            curr_image.append(curr_image[0])
            curr_image = torch.cat(curr_image, dim=0).permute((1,0,2,3,4)) # 4,2,3,240,320 the second dim is temporal
        else:
        # if len(curr_image.shape) == 5:  # 1,2,3,270,480
            curr_image = curr_image[-1].squeeze(0)

        # image_data = torch.chunk(curr_image, curr_image.shape[0], dim=0)  # left, right ,wrist
        # image_list = []
        # for each in image_data:
        #     each = cv2.resize(cv2.cvtColor(each.squeeze().permute(1,2,0).cpu().numpy(), cv2.COLOR_BGRA2BGR), (224, 224))
        #     image_list.append(torch.tensor(each).permute(2,0,1))
        # image_data = torch.stack(image_list, dim=0)
        curr_image = curr_image.to(torch.int64).unsqueeze(0)
        model_inputs = self.multimodal_processor(text=raw_lang, images=curr_image, return_tensors="pt").to(device=self.policy.device)
        model_inputs['pixel_values'] = model_inputs['pixel_values']
        data_dict = dict(states=robo_state)
        for k, v in model_inputs.items():
            data_dict[k] = v
        return data_dict