from .qwen2_vl_modules import *
from .qwen2_vl_modules import _CONFIG_FOR_DOC
from dex_vla.utils.fusion_modules import *
class DexVLA(Qwen2VLPreTrainedModel, GenerationMixin):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config):
        super().__init__(config)
        self.visual = Qwen2VisionTransformerPretrainedModel._from_config(
            config.vision_config, attn_implementation=config._attn_implementation
        )
        self.model = Qwen2VLModel(config)
        self.vocab_size = config.vocab_size

        self.padding_side = "left"  # set it to left by default, user can use setter to change padding_sides
        self.with_llm_head = config.with_llm_head
        # self.with_external_vit = config.with_external_vit
        self.with_text_fcs = config.with_text_fcs
        self.only_using_input_embeddings = config.only_using_input_embeddings
        self.using_film = config.using_film

        self.using_first_layer_hidden_states = getattr(config, "using_first_layer_hidden_states", False)
        self.llm_loss_weight = config.llm_loss_weight
        self.using_state = getattr(config, "using_state", False)

        self.external_vision_encoder = getattr(config, "external_vision_encoder", "None")

        if isinstance(config.policy_head_config, dict):
            config.policy_head_config = AutoConfig.for_model(**config.policy_head_config)
        self.policy_head = AutoModel.from_config(config=config.policy_head_config)

        # if self.with_llm_head:
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        if self.using_state:
            self.state_proj = nn.Linear(config.policy_head_config.state_dim, config.hidden_size)

        # Initialize weights and apply final processing
        self.post_init()
        if config.policy_head_config.model_type == "dit_diffusion_policy":
            self.policy_head.initialize_weights()
        self.input_action_proj = ActionProjector(config.hidden_size, config.hidden_size)
        self.reasoning_action_proj = ActionProjector(config.hidden_size, config.hidden_size)

        if self.using_film:
            self.reasoning_film = FiLM(feature_dim=config.hidden_size, condition_dim=config.hidden_size)

        if self.external_vision_encoder == 'resnet':
            from ..external_vision_encoder.resnet_vision_encoder import ResNetEncoder
            self.external_vision_encoder_model = ResNetEncoder()
            pass
        # elif self.using_xattn:
        #     # xattn_config = {
        #     #     hidden_size:
        #     # }
        #     xattn_config = SimpleNamespace(
        #         hidden_size=config.hidden_size,
        #         num_attention_heads=config.num_attention_heads,
        #         encoder_width=config.hidden_size,
        #         attention_probs_dropout_prob=0.2,
        #         hidden_dropout_prob=0.2,
        #         layer_norm_eps=1e-12,
        #     )
        #     self.xattn = BertAttention(xattn_config)
    def get_input_embeddings(self):
        return self.model.embed_tokens

    def set_input_embeddings(self, value):
        self.model.embed_tokens = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def set_decoder(self, decoder):
        self.model = decoder

    def get_decoder(self):
        return self.model

    def get_rope_index(
        self,
        input_ids: torch.LongTensor,
        image_grid_thw: Optional[torch.LongTensor] = None,
        video_grid_thw: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Calculate the 3D rope index based on image and video's temporal, height and width in LLM.

        Explanation:
            Each embedding sequence contains vision embedding and text embedding or just contains text embedding.

            For pure text embedding sequence, the rotary position embedding has no difference with mordern LLMs.
            Examples:
                input_ids: [T T T T T], here T is for text.
                temporal position_ids: [0, 1, 2, 3, 4]
                height position_ids: [0, 1, 2, 3, 4]
                width position_ids: [0, 1, 2, 3, 4]

            For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part
            and 1D rotary position embeddin for text part.
            Examples:
                Assume we have a video input with 3 temporal patches, 2 height patches and 2 width patches.
                input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.
                vision temporal position_ids: [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]
                vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]
                vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]
                text temporal position_ids: [3, 4, 5, 6, 7]
                text height position_ids: [3, 4, 5, 6, 7]
                text width position_ids: [3, 4, 5, 6, 7]
                Here we calculate the text start position_ids as the max vision position_ids plus 1.

        Args:
            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
                Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
                it.
            image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
                The temporal, height and width of feature shape of each image in LLM.
            video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
                The temporal, height and width of feature shape of each video in LLM.
            attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
                Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:

                - 1 for tokens that are **not masked**,
                - 0 for tokens that are **masked**.

        Returns:
            position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
            mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
        """
        spatial_merge_size = self.config.vision_config.spatial_merge_size
        image_token_id = self.config.image_token_id
        video_token_id = self.config.video_token_id
        vision_start_token_id = self.config.vision_start_token_id
        mrope_position_deltas = []
        if image_grid_thw is not None or video_grid_thw is not None:
            total_input_ids = input_ids
            position_ids = torch.ones(
                3, input_ids.shape[0], input_ids.shape[1], dtype=input_ids.dtype, device=input_ids.device
            )
            image_index, video_index = 0, 0
            for i, input_ids in enumerate(total_input_ids):
                if attention_mask is not None:
                    input_ids = input_ids[attention_mask[i] == 1]
                image_nums, video_nums = 0, 0
                vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)
                vision_tokens = input_ids[vision_start_indices + 1]
                image_nums = (vision_tokens == image_token_id).sum()
                video_nums = (vision_tokens == video_token_id).sum()
                input_tokens = input_ids.tolist()
                llm_pos_ids_list: list = []
                st = 0
                remain_images, remain_videos = image_nums, video_nums
                for _ in range(image_nums + video_nums):
                    if image_token_id in input_tokens and remain_images > 0:
                        ed_image = input_tokens.index(image_token_id, st)
                    else:
                        ed_image = len(input_tokens) + 1
                    if video_token_id in input_tokens and remain_videos > 0:
                        ed_video = input_tokens.index(video_token_id, st)
                    else:
                        ed_video = len(input_tokens) + 1
                    if ed_image < ed_video:
                        t, h, w = (
                            image_grid_thw[image_index][0],
                            image_grid_thw[image_index][1],
                            image_grid_thw[image_index][2],
                        )
                        image_index += 1
                        remain_images -= 1
                        ed = ed_image
                    else:
                        t, h, w = (
                            video_grid_thw[video_index][0],
                            video_grid_thw[video_index][1],
                            video_grid_thw[video_index][2],
                        )
                        video_index += 1
                        remain_videos -= 1
                        ed = ed_video
                    llm_grid_t, llm_grid_h, llm_grid_w = (
                        t.item(),
                        h.item() // spatial_merge_size,
                        w.item() // spatial_merge_size,
                    )
                    text_len = ed - st

                    st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
                    llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)

                    t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()
                    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
                    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
                    llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
                    st = ed + llm_grid_t * llm_grid_h * llm_grid_w

                if st < len(input_tokens):
                    st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
                    text_len = len(input_tokens) - st
                    llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)

                llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
                position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)
                mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))
            mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
            return position_ids, mrope_position_deltas
        else:
            if attention_mask is not None:
                position_ids = attention_mask.long().cumsum(-1) - 1
                position_ids.masked_fill_(attention_mask == 0, 1)
                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(input_ids.device)
                max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]
                mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]
            else:
                position_ids = (
                    torch.arange(input_ids.shape[1], device=input_ids.device)
                    .view(1, 1, -1)
                    .expand(3, input_ids.shape[0], -1)
                )
                mrope_position_deltas = torch.zeros(
                    [input_ids.shape[0], 1],
                    device=input_ids.device,
                    dtype=input_ids.dtype,
                )

            return position_ids, mrope_position_deltas

    def _update_model_kwargs_for_generation(
        self,
        outputs: ModelOutput,
        model_kwargs: Dict[str, Any],
        is_encoder_decoder: bool = False,
        num_new_tokens: int = 1,
    ) -> Dict[str, Any]:
        model_kwargs = super()._update_model_kwargs_for_generation(
            outputs=outputs,
            model_kwargs=model_kwargs,
            is_encoder_decoder=is_encoder_decoder,
            num_new_tokens=num_new_tokens,
        )

        if getattr(outputs, "rope_deltas", None) is not None:
            model_kwargs["rope_deltas"] = outputs.rope_deltas

        return model_kwargs

    @add_start_docstrings_to_model_forward(QWEN2_VL_INPUTS_DOCSTRING)
    @replace_return_docstrings(output_type=Qwen2VLCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        pixel_values: Optional[torch.Tensor] = None,
        pixel_values_videos: Optional[torch.FloatTensor] = None,
        image_grid_thw: Optional[torch.LongTensor] = None,
        video_grid_thw: Optional[torch.LongTensor] = None,
        rope_deltas: Optional[torch.LongTensor] = None,
        actions: Optional[torch.LongTensor] = None,
        states: Optional[torch.FloatTensor] = None,
        is_pad: bool = False,
        is_eval: bool = False,
        tinyvla: bool = False,
        raw_images: Optional[torch.Tensor] = None,
    ) -> Union[Tuple, Qwen2VLCausalLMOutputWithPast]:
        r"""
        Args:
            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

        Returns:

        Example:

        ```python
        >>> from PIL import Image
        >>> import requests
        >>> from transformers import AutoProcessor, Qwen2VLForConditionalGeneration

        >>> model = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
        >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")

        >>> messages = [
            {
                "role": "user",
                "content": [
                    {"type": "image"},
                    {"type": "text", "text": "What is shown in this image?"},
                ],
            },
        ]
        >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
        >>> image = Image.open(requests.get(url, stream=True).raw)

        >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos])

        >>> # Generate
        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
        ```"""
        # print("@@"*40)
        # for each in [attention_mask, labels, pixel_values, pixel_values_videos, image_grid_thw, actions, states]:
        #    try:
        #        print(each.dtype)
        #    except:
        #        print('None')
        # print(self.model.dtype)
        # exit(0)
        # print(actions.shape)
        self.computed_type = torch.bfloat16
        input_ids=input_ids.to("cuda")
        attention_mask=attention_mask.to("cuda")
        if not is_eval:
            labels = labels.to("cuda")
            actions = actions.to(dtype=self.computed_type, device='cuda')
            states = states.to(dtype=self.computed_type, device='cuda')

            if self.using_state:
                attn_mask_state_embedding = torch.ones((attention_mask.shape[0], 1), dtype=torch.bool, device=attention_mask.device)
                attention_mask = torch.cat([attention_mask, attn_mask_state_embedding], dim=-1)

                state_id = torch.ones((input_ids.shape[0], 1), dtype=input_ids.dtype, device=input_ids.device)
                input_ids = torch.cat([input_ids, state_id], dim=-1)
                # start_idxs = torch.where(labels != -100)
                labels = torch.cat([-100 * torch.ones((labels.shape[0], 1), dtype=labels.dtype, device=labels.device), labels], dim=1)

                mask = (input_ids == 151653).int()  # shape: [4, 100]
                start_idxs = torch.argmax(torch.flip(mask, dims=[1]), dim=1)  # shape: [4, 100]
                start_idxs = mask.shape[1] - 1 - start_idxs
            position_ids, rope_deltas = self.get_rope_index(
                input_ids, image_grid_thw, video_grid_thw, attention_mask
            )
        else:
            if self.using_state:
                mask = (input_ids == 151653).int()  # shape: [4, 100]
                start_idxs = torch.argmax(torch.flip(mask, dims=[1]), dim=1)  # shape: [4, 100]
                start_idxs = mask.shape[1] - 1 - start_idxs

        if pixel_values is not None:
            pixel_values = pixel_values.to(dtype=self.computed_type, device='cuda')
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if inputs_embeds is None:
            inputs_embeds = self.model.embed_tokens(input_ids)
            ground_truth_reasoning_embed = inputs_embeds
            if pixel_values is not None:
                pixel_values = pixel_values.type(self.visual.get_dtype())
                image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
                n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
                n_image_features = image_embeds.shape[0]
                if n_image_tokens != n_image_features:
                    raise ValueError(
                        f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
                    )
                image_mask = (
                    (input_ids == self.config.image_token_id)
                    .unsqueeze(-1)
                    .expand_as(inputs_embeds)
                    .to(inputs_embeds.device)
                )
                image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
                inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)

            if pixel_values_videos is not None:
                pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype())
                video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
                n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
                n_video_features = video_embeds.shape[0]
                if n_video_tokens != n_video_features:
                    raise ValueError(
                        f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
                    )
                video_mask = (
                    (input_ids == self.config.video_token_id)
                    .unsqueeze(-1)
                    .expand_as(inputs_embeds)
                    .to(inputs_embeds.device)
                )
                video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
                inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)

            if attention_mask is not None:
                attention_mask = attention_mask.to(inputs_embeds.device)

        if self.using_state and inputs_embeds.shape[1] > 1:
            state_embedding = self.state_proj(states).unsqueeze(1)
            temp = []
            for id, each in enumerate(inputs_embeds):
                s_id = start_idxs[id] + 1
                t = torch.cat([each[:s_id], state_embedding[id], each[s_id:-1]], dim=0)
                temp.append(t)

            inputs_embeds = torch.stack(temp, dim=0)

        outputs = self.model(
            input_ids=None,
            position_ids=position_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=True,
            return_dict=return_dict,
        )
        hidden_states = outputs[0]
        if tinyvla:
            return hidden_states
        if self.with_llm_head:
            logits = self.lm_head(hidden_states)
            logits = logits.float()
        else:
            logits = None
            self.llm_head = None

        llm_loss = None
        if labels is not None and self.with_llm_head:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            llm_loss = loss_fct(shift_logits, shift_labels)
        
        if is_eval:
            loss = None
            if not return_dict:
                # print(f"!@#@#@#@#@#$#$#####$$$$$$$@#@##@##@#@#@###############@#@#@#@@@@@@@@@@@@@@{return_dict}@@@@@@@@@@@@")
                output = (logits,) + outputs[1:]
                return (loss,) + output if loss is not None else output

            return Qwen2VLCausalLMOutputWithPast(
                loss=loss,
                logits=logits,
                past_key_values=outputs.past_key_values,
                hidden_states=outputs.hidden_states,
                attentions=outputs.attentions,
                rope_deltas=rope_deltas,
            )

        action_hidden_states = self.fusion_input_reasoning(labels=labels, input_ids=input_ids,
                                                           prev_layer_hidden_states=outputs.hidden_states,
                                                           hidden_states=hidden_states)
        # else:
        #     action_hidden_states = hidden_states
        external_obs_cond = None
        if self.external_vision_encoder == 'resnet':
            external_obs_cond = self.external_vision_encoder_model(raw_images)

        ret = self.policy_head(actions=actions, hidden_states=action_hidden_states, states=states, is_pad=is_pad, external_obs_cond=external_obs_cond)

        if self.with_llm_head:
            loss = {'loss': ret['loss'] + self.llm_loss_weight * llm_loss,
                         'llm_loss': llm_loss,
                         'action_loss': ret['loss']}
        else:
            loss = {'loss': ret['loss'],
                         'llm_loss': (torch.ones(1)*(-100)).to(ret['loss'].dtype).squeeze(0),
                         'action_loss': ret['loss']}
        if not return_dict:
            output = (logits,) + outputs[1:]
            return (loss,) + output if loss is not None else output

        torch.cuda.empty_cache()
        gc.collect()
        del input_ids
        del attention_mask
        del position_ids
        del past_key_values
        del inputs_embeds
        del labels
        del pixel_values
        del image_grid_thw
        del actions
        del states
        return Qwen2VLCausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            rope_deltas=rope_deltas,
        )

    def fusion_input_reasoning(self, labels, input_ids, hidden_states, prev_layer_hidden_states=None):
        inputs_index = labels[:, :] == -100
        inputs_index = inputs_index.int()
        # diff = inputs_index[:, :1] - inputs_index[:, 1:]
        # indexs = torch.argmax((diff!=0).float(), dim=1)
        xor_array = torch.bitwise_xor(inputs_index[:, :-1], inputs_index[:, 1:])
        indexs = torch.argmax((xor_array != 0).float(), dim=1)
        input_embeddings = []
        reasoning_embeddings = []
        identity = []
        input_embeddings_first_layer = []
        for i in range(indexs.shape[0]):
            end = indexs[i] + 1
            temp = input_ids[i] == 151643  # pad token id
            start = sum(temp.int())
            input_embeddings.append(self.input_action_proj(hidden_states[i, start:end, :]))
            identity.append(torch.mean(hidden_states[i, start:end, :], dim=0))

            input_embeddings_first_layer.append(torch.mean(prev_layer_hidden_states[0][i, start:end, :], dim=0))

            reasoning_embeddings.append(self.reasoning_action_proj(hidden_states[i, end:, :]))
        input_embeddings = torch.cat(input_embeddings, dim=0)
        input_embeddings_first_layer = torch.stack(input_embeddings_first_layer)
        reasoning_embeddings = torch.cat(reasoning_embeddings, dim=0)
        identity = torch.stack(identity)

        if self.using_film:
            action_hidden_states = self.reasoning_film(input_embeddings, reasoning_embeddings).unsqueeze(1)
            action_hidden_states = action_hidden_states + identity.unsqueeze(1)
        else:
            # action_hidden_states = input_embeddings.unsqueeze(1) + reasoning_embeddings.unsqueeze(1)
            action_hidden_states = identity.unsqueeze(1)

        if self.using_first_layer_hidden_states:
            action_hidden_states = action_hidden_states + input_embeddings_first_layer.unsqueeze(1)
        return action_hidden_states

    def xattn_forward(self, labels, input_ids, hidden_states):
        inputs_index = labels[:, :] == -100
        inputs_index = inputs_index.int()
        # diff = inputs_index[:, :1] - inputs_index[:, 1:]
        # indexs = torch.argmax((diff!=0).float(), dim=1)
        xor_array = torch.bitwise_xor(inputs_index[:, :-1], inputs_index[:, 1:])
        indexs = torch.argmax((xor_array != 0).float(), dim=1)
        input_embeddings = []
        reasoning_embeddings = []
        for i in range(indexs.shape[0]):
            end = indexs[i] + 1
            temp = input_ids[i] == 151643  # pad token id
            start = sum(temp.int())
            input_embeddings.append(torch.flip(hidden_states[i, start:end, :], dims=[1]))
            reasoning_embeddings.append(torch.flip(hidden_states[i, end:, :], dims=[1]))
        padd_value = 0

        input_embeddings = torch.nn.utils.rnn.pad_sequence(input_embeddings,
                                                           batch_first=True,
                                                           padding_value=padd_value)
        reasoning_embeddings = torch.nn.utils.rnn.pad_sequence(reasoning_embeddings,
                                                               batch_first=True,
                                                               padding_value=padd_value)

        input_embeddings = torch.flip(input_embeddings, dims=[1])
        reasoning_embeddings = torch.flip(reasoning_embeddings, dims=[1])

        r_mask = reasoning_embeddings[..., 0].ne(padd_value).float()
        i_mask = input_embeddings[..., 0].ne(padd_value).float()
        x_attn_mask = torch.matmul(r_mask.unsqueeze(-1), i_mask.unsqueeze(-1).transpose(-1, -2))
        x_attn_mask = x_attn_mask.masked_fill(x_attn_mask == 0, float(-50)).masked_fill(x_attn_mask == 1, float(0.0))
        input_embeddings = input_embeddings.to(torch.bfloat16)
        reasoning_embeddings = reasoning_embeddings.to(torch.bfloat16)
        x_attn_mask = x_attn_mask.to(torch.bfloat16)
        action_hidden_states = self.xattn(encoder_hidden_states=input_embeddings, hidden_states=reasoning_embeddings,
                                          encoder_attention_mask=x_attn_mask)[0]
        r_mask = r_mask.unsqueeze(-1).expand_as(action_hidden_states)
        action_hidden_states = action_hidden_states * r_mask
        sum_val = action_hidden_states.sum(dim=1)
        counts = r_mask.sum(dim=1).float()
        action_hidden_states = sum_val / (counts)
        action_hidden_states = action_hidden_states.unsqueeze(1).to(torch.bfloat16)
        # action_hidden_states = torch.mean(action_hidden_states, dim=1).unsqueeze(1)
        return action_hidden_states

    def prepare_inputs_for_generation(
            self,
            input_ids,
            past_key_values=None,
            attention_mask=None,
            inputs_embeds=None,
            cache_position=None,
            position_ids=None,
            use_cache=True,
            pixel_values=None,
            pixel_values_videos=None,
            image_grid_thw=None,
            video_grid_thw=None,
            **kwargs,
    ):
        # If we have cache: let's slice `input_ids` through `cache_position`, to keep only the unprocessed tokens
        # Exception 1: when passing input_embeds, input_ids may be missing entries
        # Exception 2: some generation methods do special slicing of input_ids, so we don't need to do it here
        if past_key_values is not None:
            if inputs_embeds is not None:  # Exception 1
                input_ids = input_ids[:, -cache_position.shape[0]:]
            elif input_ids.shape[1] != cache_position.shape[0]:  # Default case (the "else", a no op, is Exception 2)
                input_ids = input_ids[:, cache_position]

        rope_deltas = kwargs.get("rope_deltas", None)
        if attention_mask is not None and position_ids is None:
            if cache_position is None or (cache_position is not None and cache_position[0] == 0):
                position_ids, rope_deltas = self.get_rope_index(
                    input_ids, image_grid_thw, video_grid_thw, attention_mask
                )
            else:
                batch_size, seq_length = input_ids.shape
                delta = (
                    cache_position[0] + rope_deltas if cache_position is not None and rope_deltas is not None else 0
                )
                position_ids = torch.arange(seq_length, device=input_ids.device)
                position_ids = position_ids.view(1, -1).expand(batch_size, -1)
                position_ids = position_ids.add(delta)
                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)

        if cache_position[0] != 0:
            pixel_values = None
            pixel_values_videos = None

        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
        if inputs_embeds is not None and cache_position[0] == 0:
            model_inputs = {"inputs_embeds": inputs_embeds, "input_ids": None}
        else:
            model_inputs = {"input_ids": input_ids, "inputs_embeds": None}

        if isinstance(past_key_values, StaticCache) and attention_mask.ndim == 2:
            if model_inputs["inputs_embeds"] is not None:
                batch_size, sequence_length, _ = inputs_embeds.shape
                device = inputs_embeds.device
            else:
                batch_size, sequence_length = input_ids.shape
                device = input_ids.device

            attention_mask = self.model._prepare_4d_causal_attention_mask_with_cache_position(
                attention_mask,
                sequence_length=sequence_length,
                target_length=past_key_values.get_max_cache_shape(),
                dtype=self.lm_head.weight.dtype,
                device=device,
                cache_position=cache_position,
                batch_size=batch_size,
                config=self.config,
                past_key_values=past_key_values,
            )

        model_inputs.update(
            {
                "position_ids": position_ids,
                "past_key_values": past_key_values,
                "use_cache": use_cache,
                "attention_mask": attention_mask,
                "pixel_values": pixel_values,
                "pixel_values_videos": pixel_values_videos,
                "image_grid_thw": image_grid_thw,
                "video_grid_thw": video_grid_thw,
                "rope_deltas": rope_deltas,
            }
        )
        model_inputs.update(kwargs)
        return model_inputs

    def evaluate(self,
                 input_ids: torch.LongTensor = None,
                 actions=None,
                 states=None,
                 is_pad=None,
                 tokenizer=None,
                 is_eval=True,
                 select_one=False,
                 pixel_values=None,
                 policy_config=None,
                 attention_mask=None,
                 image_grid_thw=None,
                 raw_images=None,
                 ):
        input_ids = input_ids.to('cuda')
        if self.using_state:
            state_id = torch.ones((input_ids.shape[0], 1), dtype=input_ids.dtype, device=input_ids.device)
            input_ids = torch.cat([input_ids, state_id], dim=-1)
            attn_mask_state_embedding = torch.ones((attention_mask.shape[0], 1), dtype=torch.bool,
                                                   device=attention_mask.device)
            attention_mask = torch.cat([attention_mask, attn_mask_state_embedding], dim=-1)

        with torch.inference_mode():
            outputs = self.generate(
                input_ids,
                pixel_values=pixel_values,
                attention_mask=attention_mask,
                image_grid_thw=image_grid_thw,
                is_eval=is_eval,
                states=states.to(torch.bfloat16),
                num_beams=1,
                do_sample=False,
                temperature=0.2,
                max_new_tokens=60,
                eos_token_id=tokenizer.eos_token_id,  # End of sequence token
                pad_token_id=tokenizer.eos_token_id,  # Pad token
                use_cache=True,
                output_hidden_states=True,
                return_dict_in_generate=True,
            )

        output_ids = outputs.sequences
        # last_hidden_states = outputs.hidden_states[-2][-1]
        input_token_len = input_ids.shape[1]
        n_diff_input_output = (input_ids != output_ids[:, :input_token_len]).sum().item()
        if n_diff_input_output > 0:
            print(f'[Warning] {n_diff_input_output} output_ids are not the same as the input_ids')
        outputs_text = tokenizer.batch_decode(output_ids[:, input_token_len:], skip_special_tokens=False)[0]

        outputs_text = outputs_text.strip()
        new_output_ids = tokenizer.encode(outputs_text)
        last_hidden_states = [each[-1] for each in outputs.hidden_states]  # all hidden states
        all_hidden_states = torch.cat(last_hidden_states, dim=1)
        ground_truth_reasoning_embed = self.model.embed_tokens(torch.tensor(new_output_ids).to('cuda')).unsqueeze(0)
        new_output_ids = torch.tensor([new_output_ids]).to(device=input_ids.device)
        labels = torch.cat([torch.ones_like(input_ids) * -100, new_output_ids], dim=-1)

        action_hidden_states = None
        external_obs_cond = None
        if self.external_vision_encoder == 'resnet':
            external_obs_cond = self.external_vision_encoder_model(raw_images)
        action_hidden_states = self.fusion_input_reasoning(labels=labels,
                                                       input_ids=torch.cat([input_ids, new_output_ids], dim=-1),
                                                       hidden_states=torch.cat(last_hidden_states, dim=1),
                                                       prev_layer_hidden_states=outputs.hidden_states[0])


        # print(outputs)
        action = self.policy_head(actions, action_hidden_states, states.to(all_hidden_states.dtype), is_pad, external_obs_cond=external_obs_cond)
        return action, outputs_text

    def evaluate_tinyvla(self,
                 input_ids: torch.LongTensor = None,
                 actions=None,
                 states=None,
                 is_pad=None,
                 tokenizer=None,
                 is_eval=True,
                 select_one=False,
                 pixel_values=None,
                 policy_config=None,
                 attention_mask=None,
                 image_grid_thw=None,
                 ):
        input_ids = input_ids.to('cuda')
        with torch.inference_mode():
            all_hidden_states = self.forward(input_ids,
                pixel_values=pixel_values,
                attention_mask=attention_mask,
                image_grid_thw=image_grid_thw,
                is_eval=is_eval,
                tinyvla=True)

        # last_hidden_states = outputs.hidden_states[-2][-1]
        # select_one = False
        all_hidden_states = torch.mean(all_hidden_states, dim=1).unsqueeze(1)

        # print(outputs)
        action = self.policy_head(actions, all_hidden_states, states.to(all_hidden_states.dtype), is_pad)
        return action, "tinyvla no output"

from transformers import AutoModelForCausalLM
AutoModelForCausalLM.register(DexVLAConfig, DexVLA)
