from typing import Optional, List

import torch

from .modeling_internvl_chat import InternVLChatModel
from .configuration_tinyvla import TinyVLAConfig

from transformers import AutoConfig, AutoModel

class TinyVLA(InternVLChatModel):
    config_class = TinyVLAConfig
    def __init__(self, config: TinyVLAConfig):
        super().__init__(config)
        self.system_message = "You are a helpful assistant."
        # remove lm_head for tinyvla
        self.language_model.lm_head = None

        # setup for policy head
        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)
        self.post_init()

    def forward(
            self,
            pixel_values: torch.FloatTensor,
            actions: torch.FloatTensor,
            states: torch.FloatTensor,
            input_ids: torch.LongTensor = None,
            attention_mask: Optional[torch.Tensor] = None,
            position_ids: Optional[torch.LongTensor] = None,
            image_flags: Optional[torch.LongTensor] = None,
            past_key_values: Optional[List[torch.FloatTensor]] = None,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
            is_pad: Optional[bool] = None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        hidden_states = self.vlm_forward(
            pixel_values=pixel_values,
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            return_dict=return_dict
        )

        loss = self.policy_head_forward(hidden_states, actions, states, is_pad)
        return (loss, ) + (hidden_states, )


    def vlm_forward(
            self,
            pixel_values: torch.FloatTensor,
            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,
            use_cache: Optional[bool] = None,
            output_attentions: Optional[bool] = None,
            output_hidden_states: Optional[bool] = None,
            return_dict: Optional[bool] = None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        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
        )

        input_embeds = self.language_model.get_input_embeddings()(input_ids).clone()

        if pixel_values.ndim == 5:
            c, h, w = pixel_values.size()[2:]
            pixel_values = pixel_values.view(-1, c, h, w)
        vit_embeds = self.extract_feature(pixel_values)
        vit_batch_size = pixel_values.shape[0]

        B, N, C = input_embeds.shape
        input_embeds = input_embeds.reshape(B * N, C)

        # if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
        #     print(
        #         f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')

        input_ids = input_ids.reshape(B * N)
        selected = (input_ids == self.img_context_token_id)
        try:
            input_embeds[selected] = input_embeds[selected] * 0.0 + vit_embeds.reshape(-1, C)
        except Exception as e:
            vit_embeds = vit_embeds.reshape(-1, C)
            print(f'warning: {e}, input_embeds[selected].shape={input_embeds[selected].shape}, '
                  f'vit_embeds.shape={vit_embeds.shape}')
            n_token = min(selected.sum(), vit_embeds.size(0))
            input_embeds[selected][:n_token] = input_embeds[selected][:n_token] * 0.0 + vit_embeds[:n_token]

        input_embeds = input_embeds.reshape(B, N, C)

        # CausalLM forward
        outputs = self.language_model.model(
            inputs_embeds=input_embeds,
            attention_mask=attention_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        hidden_states = outputs[0]
        return hidden_states

    def policy_head_forward(self, hidden_states, noise_actions, states, is_pad):
        ret = self.policy_head(
            actions=noise_actions,
            hidden_states=hidden_states,
            states=states,
            is_pad=is_pad
        )
        return ret['loss']

    def sample_action(self,
                      input_ids: torch.LongTensor = None,
                      actions=None,
                      states=None,
                      is_pad=None,
                      pixel_values=None,
                      attention_mask=None,
                      ):
        input_ids = input_ids.to("cuda")
        states = states.to("cuda")
        pixel_values = pixel_values.to("cuda")
        attention_mask = attention_mask.to("cuda")

        hidden_states = self.vlm_forward(
            pixel_values=pixel_values,
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=None
        )

        all_hidden_states = torch.mean(hidden_states, dim=1).unsqueeze(1)
        action = self.policy_head(actions, all_hidden_states, states.to(all_hidden_states.dtype), is_pad)
        return action

from transformers import AutoModelForCausalLM
AutoModelForCausalLM.register(TinyVLAConfig, TinyVLA)