import torch

from PIL import Image
from qwen_vl_utils import fetch_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 qwen2_vla_policy:
    def __init__(self, policy_config, data_args=None):
        super(qwen2_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]"]})
        
        paths = model_path.split('/')[:-1]
        if 'checkpoint' in paths[-1]:
            paths = paths[:-1]
        self.config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
    def datastruct_droid2qwen2vla(self, raw_lang, image_len):
        messages = [
            {
                "role": "user",
                "content": [],
            },
            # {"role": "assistant", "content": f''},
        ]

        for i in range(image_len):
            messages[0]['content'].append({
                        "type": "image",
                        "image": None,
                    })

        messages[0]['content'].append({"type": "text", "text": f""})

        messages[0]['content'][-1]['text'] = raw_lang
        # messages[1]['content'] = sample['reasoning'] + "Next action:"
        # print(sample['obs']['raw_language'].decode('utf-8'))
        return messages

    def qwen2_image_preprocess(self, each, camera_name):
        ele = {
            # "resized_height": None,
            # "resized_width": None
        }
        each = Image.fromarray(each.squeeze(0).permute(1, 2, 0).cpu().numpy().astype(np.uint8))
        ele['image'] = each
        # if 'wrist' in camera_name:
        #     # w, h = eval(self.data_args.image_size_wrist)
        #     w,h=224,224
        #     ele['resized_height'] = h
        #     ele['resized_width'] = w
        # else:
        #     ele['resized_height'] = each.height
        #     ele['resized_width'] = each.width
        ele['resized_height'] = each.height
        ele['resized_width'] = each.width
        each = fetch_image(ele)
        return torch.from_numpy(np.array(each))

    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)

        messages = self.datastruct_droid2qwen2vla(raw_lang, curr_image.shape[0])
        image_data = torch.chunk(curr_image, curr_image.shape[0], dim=0)  # left, right ,wrist
        image_list = []
        for i, each in enumerate(image_data[:]):
            each = each.squeeze(0)
            if each.ndim == 3:
                img_pil = self.qwen2_image_preprocess(each, CAMERA_VIEWS[i])
            else:
                img_pil = []
                for temp in each.squeeze(0):
                    img_pil.append(self.qwen2_image_preprocess(temp, CAMERA_VIEWS[i]))
                img_pil = torch.stack(img_pil, 0)
            image_list.append(img_pil)

        # TODO RESIZE
        # image_data = image_data / 255.0
        image_data = image_list
        text = self.multimodal_processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        # image_inputs, video_inputs = process_vision_info(dataset)
        # text = text[:-23]
        video_inputs = None
        model_inputs = self.multimodal_processor(
            text=text,
            images=image_data,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        data_dict = dict(states=robo_state)
        for k, v in model_inputs.items():
            data_dict[k] = v
        return data_dict