from typing import Optional, List
import torch
from torch import nn
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
from transformers import PreTrainedModel, AutoModel, AutoModelForCausalLM
from transformers.modeling_outputs import ModelOutput
from transformers.generation.utils import GenerationMixin
from transformers.cache_utils import Cache, DynamicCache
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
from transformers.models.qwen3.modeling_qwen3 import eager_attention_forward, BaseModelOutputWithPast

from .modeling_moonvit import patch_merger, get_rope_index, apply_multimodal_rotary_pos_emb
from .configuration_smallvlm import SmallVLMConfig

class Qwen2_5_VLRotaryEmbedding(nn.Module):
    def __init__(self, config, device=None):
        super().__init__()
        # BC: "rope_type" was originally "type"
        if hasattr(config, "rope_scaling") and config.rope_scaling is not None:
            self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
        else:
            self.rope_type = "default"
        self.max_seq_len_cached = config.max_position_embeddings
        self.original_max_seq_len = config.max_position_embeddings

        self.config = config
        self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]

        inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.original_inv_freq = self.inv_freq

    @torch.no_grad()
    @dynamic_rope_update  # power user: used with advanced RoPE types (e.g. dynamic rope)
    def forward(self, x, position_ids):
        # In contrast to other models, Qwen2_5_VL has different position ids for the grids
        # So we expand the inv_freq to shape (3, ...)
        inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
        position_ids_expanded = position_ids[:, :, None, :].float()  # shape (3, bs, 1, positions)

        device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
        with torch.autocast(device_type=device_type, enabled=False):  # Force float32
            freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
            emb = torch.cat((freqs, freqs), dim=-1)
            cos = emb.cos() * self.attention_scaling
            sin = emb.sin() * self.attention_scaling

        return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
    

def build_vision_model(config, model=None):
    if model is None:
        model = AutoModel.from_config(config, trust_remote_code=True)
    return model

def mrope_forward(
    self,
    hidden_states: torch.Tensor,
    position_embeddings: Tuple[torch.Tensor, torch.Tensor],
    attention_mask: Optional[torch.Tensor],
    past_key_value: Optional[Cache] = None,
    cache_position: Optional[torch.LongTensor] = None,
    **kwargs,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
    input_shape = hidden_states.shape[:-1]
    hidden_shape = (*input_shape, -1, self.head_dim)

    query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
    key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

    cos, sin = position_embeddings
    query_states, key_states = apply_multimodal_rotary_pos_emb(query_states, key_states, cos, sin, [16, 24, 24], unsqueeze_dim=1)

    if past_key_value is not None:
        # sin and cos are specific to RoPE models; cache_position needed for the static cache
        cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
    attention_interface: Callable = eager_attention_forward
    if self.config._attn_implementation != "eager":
        if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):
            pass
        else:
            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]

    attn_output, attn_weights = attention_interface(
        self,
        query_states,
        key_states,
        value_states,
        attention_mask,
        dropout=0.0 if not self.training else self.attention_dropout,
        scaling=self.scaling,
        sliding_window=self.sliding_window,  # diff with Llama
        **kwargs,
    )

    attn_output = attn_output.reshape(*input_shape, -1).contiguous()
    attn_output = self.o_proj(attn_output)
    return attn_output, attn_weights

import transformers
transformers.models.qwen3.modeling_qwen3.Qwen3Attention.forward = mrope_forward


def forward(
    self,
    input_ids: Optional[torch.LongTensor] = None,
    attention_mask: Optional[torch.Tensor] = None,
    position_ids: Optional[torch.LongTensor] = None,
    past_key_values: Optional[Cache] = None,
    inputs_embeds: Optional[torch.FloatTensor] = None,
    use_cache: Optional[bool] = None,
    output_attentions: Optional[bool] = None,
    output_hidden_states: Optional[bool] = None,
    cache_position: Optional[torch.LongTensor] = None,
    **flash_attn_kwargs,
) -> BaseModelOutputWithPast:
    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
    )
    use_cache = use_cache if use_cache is not None else self.config.use_cache

    if (input_ids is None) ^ (inputs_embeds is not None):
        raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

    if self.gradient_checkpointing and self.training and use_cache:
        use_cache = False

    # TODO (joao): remove this exception in v4.56 -- it exists for users that try to pass a legacy cache
    if not isinstance(past_key_values, (type(None), Cache)):
        raise ValueError("The `past_key_values` should be either a `Cache` object or `None`.")

    if inputs_embeds is None:
        inputs_embeds = self.embed_tokens(input_ids)

    if use_cache and past_key_values is None:
        past_key_values = DynamicCache()

    if cache_position is None:
        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
        cache_position = torch.arange(
            past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
        )

    if position_ids is None:
        position_ids = cache_position.unsqueeze(0)

    causal_mask = self._update_causal_mask(
        attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
    )

    hidden_states = inputs_embeds

    # create position embeddings to be shared across the decoder layers
    position_embeddings = self.rotary_emb(hidden_states, position_ids)

    # decoder layers
    all_hidden_states = () if output_hidden_states else None
    all_self_attns = () if output_attentions else None

    for decoder_layer in self.layers[: self.config.num_hidden_layers]:
        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        if self.gradient_checkpointing and self.training:
            layer_outputs = self._gradient_checkpointing_func(
                decoder_layer.__call__,
                hidden_states,
                causal_mask,
                position_ids,
                past_key_values,
                output_attentions,
                use_cache,
                cache_position,
                position_embeddings,
            )
        else:
            layer_outputs = decoder_layer(
                hidden_states,
                attention_mask=causal_mask,
                position_ids=position_ids,
                past_key_value=past_key_values,
                output_attentions=output_attentions,
                use_cache=use_cache,
                cache_position=cache_position,
                position_embeddings=position_embeddings,
                **flash_attn_kwargs,
            )

        hidden_states = layer_outputs[0]

        if output_attentions:
            all_self_attns += (layer_outputs[1],)

    hidden_states = self.norm(hidden_states)

    # add hidden states from the last decoder layer
    if output_hidden_states:
        all_hidden_states += (hidden_states,)

    return BaseModelOutputWithPast(
        last_hidden_state=hidden_states,
        past_key_values=past_key_values if use_cache else None,
        hidden_states=all_hidden_states,
        attentions=all_self_attns,
    )
transformers.models.qwen3.modeling_qwen3.Qwen3Model.forward = forward

class SmallVLMForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = SmallVLMConfig
    supports_gradient_checkpointing = True
    _skip_keys_device_placement = "past_key_values"
    _supports_cache_class = True
    _supports_flash_attn_2 = True
    _supports_sdpa = True

    def __init__(self, config, language_model=None, vision_model=None):
        super().__init__(config)
        self.rope_deltas = None  # cache rope_deltas here
        
        vision_model = build_vision_model(config.vision_model_config, vision_model)
        if language_model is None:
            kwargs_ = {}
            if config._attn_implementation_internal is not None:
                kwargs_['attn_implementation'] = config._attn_implementation_internal
            language_model = AutoModelForCausalLM.from_config(config.language_model_config, trust_remote_code=True, **kwargs_)

        self.vision_model = vision_model
        
        self.language_model = language_model

        self.vision_to_text_proj = nn.Sequential( # map the text embeddings to vision encoder
            nn.Linear(self.config.vision_model_config.hidden_size, self.config.language_model_config.hidden_size),
            nn.GELU(),
            nn.Linear(self.config.language_model_config.hidden_size, self.config.language_model_config.hidden_size)
        )

        self.text_to_vision_proj = nn.Sequential(
            nn.Linear(self.config.language_model_config.hidden_size, self.config.vision_model_config.hidden_size),
            nn.GELU(),
            nn.Linear(self.config.vision_model_config.hidden_size, self.config.vision_model_config.hidden_size)
        )
        self.vision_rotary_emb = Qwen2_5_VLRotaryEmbedding(config.vision_model_config)
        self.text_rotary_emb = Qwen2_5_VLRotaryEmbedding(config.language_model_config)
        self.language_model.model.rotary_emb = self.text_rotary_emb

        for layer in self.language_model.model.layers:
            setattr(layer.self_attn, 'layer_idx', layer.self_attn.layer_idx + self.vision_model.config.num_hidden_layers)

        self.gradient_checkpointing = False

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = 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,
        cache_position: Optional[torch.LongTensor] = None,
        pixel_values: Optional[torch.FloatTensor] = None,
        grid_hws: Optional[torch.LongTensor] = None,
    ):

        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 (input_ids is None) ^ (inputs_embeds is not None):
            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")

        if self.gradient_checkpointing and self.training and use_cache:
            use_cache = False

        if use_cache and past_key_values is None:
            past_key_values = DynamicCache()

        inputs_embeds = self.get_input_embeddings()(input_ids)
        inputs_embeds = self.text_to_vision_proj(inputs_embeds)
        is_dummy_input = pixel_values is not None and pixel_values.size(0) == 0
        if is_dummy_input:
            pixel_values = torch.zeros((4,) + pixel_values.shape[1:], dtype=pixel_values.dtype, device=pixel_values.device)
            grid_hws = torch.tensor([[1, 2, 2]], dtype=torch.int32).to(pixel_values.device)

        if pixel_values is not None:
            vision_embeds = self.vision_model.patch_embed(pixel_values, grid_hws[:, 1:])
            vision_embeds_list = patch_merger(
                vision_embeds, grid_hws[:, 1:], merge_kernel_size=self.vision_model.merge_kernel_size
            )
            vision_embeds = self.vision_model.pixel_merger(torch.cat(vision_embeds_list).view(-1, vision_embeds.shape[-1] * 4))
            
            vision_mask = (input_ids == self.config.image_token_id).to(inputs_embeds.device)
            inputs_embeds[vision_mask] = vision_embeds

            image_token_lens = (grid_hws.prod(dim=1) // 4)
            bsz, src_len = attention_mask.size()
            causal_mask = attention_mask[:, None, None, :].expand(bsz, 1, src_len, src_len).to(inputs_embeds.dtype)
            causal_mask.tril_()
            idx = 0
            for i, _ in enumerate(causal_mask):
                vision_mask = input_ids[i] == self.config.image_token_id
                while (vision_mask.sum() > 0):
                    start = torch.nonzero(vision_mask)[0][0]
                    num = image_token_lens[idx]
                    idx += 1
                    causal_mask[i, 0, start:start+num, start:start+num] = 1
                    vision_mask[start:start+num] = 0

            causal_mask = 1.0 - causal_mask
            causal_mask = causal_mask.masked_fill(causal_mask.to(torch.bool), torch.finfo(vision_embeds.dtype).min)
        else:
            causal_mask = None

        if self.is_gradient_checkpointing and torch.is_grad_enabled() and self.training:
            inputs_embeds.requires_grad_(True)

        if cache_position is None:
            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
            cache_position = torch.arange(
                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
            )

        if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
            # calculate RoPE index once per generation in the pre-fill stage only
            if (
                (cache_position is not None and cache_position[0] == 0)
                or self.rope_deltas is None
                or (past_key_values is None or past_key_values.get_seq_length() == 0)
            ):
                position_ids, rope_deltas = get_rope_index(
                    self.config.image_token_id,
                    self.config.video_token_id,
                    self.config.vision_start_token_id, 
                    spatial_merge_size=2,
                    input_ids=input_ids,
                    image_grid_thw=grid_hws,
                    video_grid_thw=None,
                    attention_mask=attention_mask
                )
                self.rope_deltas = rope_deltas
            # then use the prev pre-calculated rope-deltas to get the correct position ids
            else:
                batch_size, seq_length, _ = inputs_embeds.shape
                delta = (
                    (cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
                    if cache_position is not None
                    else 0
                )
                position_ids = torch.arange(seq_length, device=inputs_embeds.device)
                position_ids = position_ids.view(1, -1).expand(batch_size, -1)
                if cache_position is not None:  # otherwise `deltas` is an int `0`
                    delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)
                position_ids = position_ids.add(delta)
                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)

        position_embeddings = self.vision_rotary_emb(inputs_embeds, position_ids)
        inputs_embeds = self.vision_model.encoder(inputs_embeds, causal_mask, position_embeddings, past_key_values)

        # return ModelOutput(
        #     last_hidden_state=self.vision_model.projector(inputs_embeds),
        #     text_hidden_state=self.vision_to_text_proj(inputs_embeds),
        # )
    
        inputs_embeds = self.vision_to_text_proj(inputs_embeds)

        outputs = self.language_model(
            input_ids=None,
            labels=labels,
            attention_mask=causal_mask,
            position_ids=position_ids,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            cache_position=cache_position,
            return_dict=True,
        )

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

    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):
        super().gradient_checkpointing_enable(gradient_checkpointing_kwargs)
        self.language_model.enable_input_require_grads()

    def get_input_embeddings(self):
        return self.language_model.get_input_embeddings()

    def set_input_embeddings(self, value):
        self.language_model.set_input_embeddings(value)

    def get_output_embeddings(self):
        return self.language_model.get_output_embeddings()

    def set_output_embeddings(self, new_embeddings):
        self.language_model.set_output_embeddings(new_embeddings)

    def set_decoder(self, decoder):
        self.language_model.set_decoder(decoder)

    def get_decoder(self):
        return self.language_model.get_decoder()

    def tie_weights(self):
        return self.language_model.tie_weights()

    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,
        **kwargs,
    ):
        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
        model_inputs = super().prepare_inputs_for_generation(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            cache_position=cache_position,
            position_ids=position_ids,
            pixel_values=pixel_values,
            use_cache=use_cache,
            **kwargs,
        )

        # Qwen2-5-VL position_ids are prepareed with rope_deltas in forward
        model_inputs["position_ids"] = None
        if cache_position[0] != 0:
            model_inputs["pixel_values"] = None

        return model_inputs