import os.path

from torchvision import transforms
from aloha_scripts.utils import *
import time
from data_utils.dataset import set_seed
from einops import rearrange
from transformers import AutoConfig, AutoProcessor, AutoTokenizer

import sys
from policy.DexVLA.policy_heads.models.transformer_diffusion.modeling_dit_diffusion import *
from policy.DexVLA.policy_heads.models.transformer_diffusion.configuration_dit_diffusion import *
from policy.DexVLA.dex_vla.utils.image_processing_qwen2_vla import *
# from paligemma_vla.utils.processing_paligemma_vla import *
# from policy.DexVLA.dex_vla.utils.processing_qwen2_vla import *
from policy.DexVLA.evaluate.vla_policy import *
from policy.DexVLA.dex_vla.model_load_utils import load_model_for_eval
from PIL import Image
import pickle
import cv2
import numpy as np

def pre_process(robot_state_value, key, stats):
    tmp = robot_state_value
    tmp = (tmp - stats[key + '_mean']) / stats[key + '_std']
    return tmp

def get_obs(deplot_env_obs, stats, camera_views=3):
    '''
    预处理输入到DexVLA的数据
    '''
    cur_top_rgb = deplot_env_obs['cam_high']
    cur_left_rgb = deplot_env_obs['cam_left']
    cur_right_rgb = deplot_env_obs['cam_right']
    # Change the data format donot Change the robot doudong action!
    cur_top_rgb = cv2.cvtColor(cur_top_rgb, cv2.COLOR_BGRA2BGR)[:, :, ::-1]
    cur_left_rgb = cv2.cvtColor(cur_left_rgb, cv2.COLOR_BGRA2BGR)[:, :, ::-1]
    cur_right_rgb = cv2.cvtColor(cur_right_rgb, cv2.COLOR_BGRA2BGR)[:, :, ::-1]

    cur_joint_positions = deplot_env_obs['qpos']

    cur_state_np = pre_process(cur_joint_positions, 'qpos', stats)

    cur_state = cur_state_np  # deplot_env_obs['state']
    cur_state = np.expand_dims(cur_state, axis=0)

    # [2, 1, 128, 128, 3]
    # [2, 480, 480, 3]
    traj_rgb_np = np.array([cur_top_rgb, cur_left_rgb, cur_right_rgb])

    traj_rgb_np = np.expand_dims(traj_rgb_np, axis=1)
    traj_rgb_np = np.transpose(traj_rgb_np, (1, 0, 4, 2, 3))

    # print("#" * 50)
    # print(traj_rgb_np.shape)

    return cur_joint_positions, cur_state, traj_rgb_np

class qwen2_vla_policy:
    def __init__(self, policy_config, data_args=None):
        super(qwen2_vla_policy).__init__()
        self.load_policy(policy_config)
        self.data_args = data_args

    def load_policy(self, policy_config):
        self.policy_config = policy_config
        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 datastruct_droid2qwen2vla(self, raw_lang):
        messages = [
            {
                "role": "user",
                "content": [
                    {
                        "type": "image",
                        "image": None,
                    },
                    {
                        "type": "image",
                        "image": None,
                    },
                    {
                        "type": "image",
                        "image": None,
                    },
                    {"type": "text", "text": f""},
                ],
            },
            # {"role": "assistant", "content": f''},
        ]

        messages[0]['content'][-1]['text'] = raw_lang

        return messages
    def process_batch_to_qwen2_vla(self, curr_image, robo_state, raw_lang):

        if len(curr_image.shape) == 5:  # 1,2,3,270,480
            curr_image = curr_image.squeeze(0)

        messages = self.datastruct_droid2qwen2vla(raw_lang)
        image_data = torch.chunk(curr_image, curr_image.shape[0], dim=0)  # top, left_wrist, right_wrist
        image_list = []
        for i, each in enumerate(image_data):
            ele = {}
            each = Image.fromarray(each.cpu().squeeze(0).permute(1, 2, 0).numpy().astype(np.uint8))
            ele['image'] = each
            ele['resized_height'] = 240
            ele['resized_width'] = 320

            image_list.append(torch.from_numpy(np.array(each)))
        # image_data = image_data / 255.0
        image_data = image_list
        text = self.multimodal_processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )

        video_inputs = None
        model_inputs = self.multimodal_processor(
            text=text,
            images=image_data,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        ) # After this code print the tensor image tensor  dtype=torch.uint8)], 'videos': None} not recognized.
        data_dict = dict(states=robo_state)
        for k, v in model_inputs.items():
            data_dict[k] = v
        return data_dict


class DexVLA:
    def __init__(self, policy_config):
        super(DexVLA).__init__()
        self.policy_config = policy_config
        self.policy = qwen2_vla_policy(policy_config)
        print("Model Path: ", policy_config["model_path"])
        # self.image_list = []
    def get_action(self, obs):
        # image_list = self.image_list
        policy = self.policy
        raw_lang = obs["raw_lang"]
        policy_config = policy.policy_config
        # eval 模式
        policy.policy.eval()
        # 训练模型路径
        paths =policy_config["model_path"].split('/')[:-1]
        # 存储统计信息。
        stats_path = policy_config["stats_path"]
        with open(stats_path, 'rb') as f:
            stats = pickle.load(f)
        post_process = lambda a: ((a + 1) / 2) * (stats['action_max'] - stats['action_min']) + stats['action_min']
        ### Coding here
        with torch.inference_mode():
            cur_state_np_raw, robot_state, traj_rgb_np = get_obs(obs, stats, camera_views=policy_config['camera_views'])
            # 转换数据格式
            robot_state = torch.from_numpy(robot_state).float().cuda()
            curr_image = torch.from_numpy(traj_rgb_np).float().cuda()
            # print('rand crop resize is used!')
            original_size = curr_image.shape[-2:]
            ratio = 0.95
            curr_image = curr_image[...,
                         int(original_size[0] * (1 - ratio) / 2): int(original_size[0] * (1 + ratio) / 2),
                         int(original_size[1] * (1 - ratio) / 2): int(original_size[1] * (1 + ratio) / 2)]
            curr_image = curr_image.squeeze(0)
            resize_transform = transforms.Resize((240, 320), antialias=True)
            curr_image = resize_transform(curr_image)
            curr_image = curr_image.unsqueeze(0)
            # image_list.append(curr_image)
            batch = policy.process_batch_to_qwen2_vla(curr_image, robot_state, raw_lang)
            all_actions, outputs = policy.policy.evaluate(**batch, is_eval=True, tokenizer=policy.tokenizer,
                                                          raw_images=curr_image)
            # print(f"Output: {outputs}")
            actions = all_actions.squeeze(0).cpu().to(dtype=torch.float32).numpy()
            actions = post_process(actions)
            # actions = actions[:25]
            return actions


def compress_single_image(img):
    """
    将单张图片压缩为 JPEG 格式的二进制数据。

    参数:
        img (numpy.ndarray): 输入的图像数组，形状为 (height, width, channels)。

    返回:
        compressed_data (bytes): 压缩后的二进制数据。
    """
    # 使用 OpenCV 的 imencode 函数将图像编码为 JPEG 格式
    success, encoded_image = cv2.imencode('.jpg', img)
    if not success:
        raise ValueError("Failed to encode image")

    # 获取编码后的二进制数据
    compressed_data = encoded_image.tobytes()

    return compressed_data

def encode_obs(observation):  # Post-Process Observation
    """
    Process input data for VLA model。
    """
    obs = observation
    cam_high = obs["observation"]["head_camera"]["rgb"]
    cam_left = obs["observation"]["left_camera"]["rgb"]
    cam_right = obs["observation"]["right_camera"]["rgb"]
    cam_high = compress_single_image(cam_high)
    cam_left = compress_single_image(cam_left)
    cam_right = compress_single_image(cam_right)
    cam_high = cv2.imdecode(np.frombuffer(cam_high, np.uint8), cv2.IMREAD_COLOR)
    cam_left = cv2.imdecode(np.frombuffer(cam_left, np.uint8), cv2.IMREAD_COLOR)
    cam_right = cv2.imdecode(np.frombuffer(cam_right, np.uint8), cv2.IMREAD_COLOR)
    qpos = (observation["joint_action"]["left_arm"] + [observation["joint_action"]["left_gripper"]] +
            observation["joint_action"]["right_arm"] + [observation["joint_action"]["right_gripper"]])
    qpos = np.array(qpos)
    return {
        "cam_high": cam_high,
        "cam_left": cam_left,
        "cam_right": cam_right,
        "qpos": qpos,
    }


def get_model(usr_args): # from deploy_policy.yml and eval.sh (overrides)
    policy_config = {
        #### 1. Specify path to trained DexVLA(Required)#############################
        "model_path": usr_args["model_path"],
        "stats_path": usr_args["stats_path"],
        "model_base": None, # only use for lora finetune
        "enable_lora": False, # only use for lora finetune
        "action_head": usr_args['action_head'],
        "tinyvla": False,
        "camera_views": 3,
        "save_model": False,
    }
    model = DexVLA(policy_config)
    return model  # return your policy model


def eval(TASK_ENV, model, observation):
    """
    All the function interfaces below are just examples
    You can modify them according to your implementation
    But we strongly recommend keeping the code logic unchanged
    """
    obs = encode_obs(observation)  # Post-Process Observation
    instruction = TASK_ENV.get_instruction()
    obs.update({"raw_lang": str(instruction)})
    actions = model.get_action(obs)  # Get Action according to observation chunk

    for action in actions[:25]:  # Execute each step of the action
        # TASK_ENV.take_one_step_action(action)
        TASK_ENV.take_action(action)
        observation = TASK_ENV.get_obs()


def reset_model(model):  # Clean the model cache at the beginning of every evaluation episode, such as the observation window
    pass
