# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
# Copyright 2026 Tensor Auto Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
PaliGemma with Expert Module.

This module implements the PaliGemma model with an additional expert module,
specifically designed for the Pi05 policy. It combines a pre-trained PaliGemma
Vision-Language Model (VLM) with a Gemma-based expert model to handle
action generation and conditioning.
"""

import torch
from torch import nn
from transformers import (
    AutoConfig,
    Cache,
    GemmaForCausalLM,
    PaliGemmaForConditionalGeneration,
    PretrainedConfig,
    PreTrainedModel,
)
from transformers.models.auto import CONFIG_MAPPING
from transformers.models.gemma import modeling_gemma


def _preferred_dtype():
    return torch.float32 if torch.onnx.is_in_onnx_export() else torch.bfloat16


def apply_rope(x: torch.Tensor, positions: torch.Tensor, max_wavelength: int = 10_000) -> torch.Tensor:
    """Applies RoPE positions to the input tensor.

    Args:
        x: Input tensor of shape [B, L, H, D].
        positions: Position tensor of shape [B, L].
        max_wavelength: Maximum wavelength for RoPE. Defaults to 10_000.

    Returns:
        Tensor: The input tensor with RoPE applied, of shape [B, L, H, D].
    """
    d_half = x.shape[-1] // 2
    device = x.device
    dtype = x.dtype
    x = x.to(torch.float32)

    freq_exponents = (2.0 / x.shape[-1]) * torch.arange(d_half, dtype=torch.float32, device=device)
    timescale = max_wavelength**freq_exponents
    radians = positions[..., None].to(torch.float32) / timescale[None, None, :].to(torch.float32)

    radians = radians[..., None, :]

    sin = torch.sin(radians)  # .to(dtype=dtype)
    cos = torch.cos(radians)  # .to(dtype=dtype)

    x1, x2 = x.split(d_half, dim=-1)
    res = torch.empty_like(x)
    res[..., :d_half] = x1 * cos - x2 * sin
    res[..., d_half:] = x2 * cos + x1 * sin

    return res.to(dtype)


class PaliGemmaWithExpertConfig(PretrainedConfig):
    """Configuration class for PaliGemmaWithExpertModel."""

    model_type = "PaliGemmaWithExpertModel"
    sub_configs = {"paligemma_config": AutoConfig, "gemma_expert_config": AutoConfig}

    def __init__(
        self,
        paligemma_config: dict | None = None,
        gemma_expert_config: dict | None = None,
        freeze_vision_encoder: bool = True,
        train_expert_only: bool = True,
        attention_implementation: str = "eager",
        load_pretrained_paligemma: bool = False,
        discrete_action_vocab_size: int | None = None,
        dropout: float = 0.1,
        **kwargs,
    ):
        """Initializes the configuration.

        Args:
            paligemma_config: Configuration dictionary for the PaliGemma model.
            gemma_expert_config: Configuration dictionary for the Gemma expert model.
            freeze_vision_encoder: Whether to freeze the vision encoder. Defaults to True.
            train_expert_only: Whether to train only the expert model. Defaults to True.
            attention_implementation: Attention implementation to use ("eager" or "fa2"). Defaults to "eager".
            load_pretrained_paligemma: Whether to load a pretrained PaliGemma model. Defaults to False.
            discrete_action_vocab_size: Vocabulary size for discrete actions.
            dropout: Dropout probability. Defaults to 0.1.
            **kwargs: Additional keyword arguments passed to PretrainedConfig.
        """
        self.freeze_vision_encoder = freeze_vision_encoder
        self.train_expert_only = train_expert_only
        self.attention_implementation = attention_implementation
        self.load_pretrained_paligemma = load_pretrained_paligemma
        self.discrete_action_vocab_size = discrete_action_vocab_size
        self.dropout = dropout

        if paligemma_config is None:
            # Default config from Pi0
            self.paligemma_config = CONFIG_MAPPING["paligemma"](
                transformers_version="4.48.1",
                _vocab_size=257152,
                bos_token_id=2,
                eos_token_id=1,
                hidden_size=2048,
                image_token_index=257152,
                model_type="paligemma",
                pad_token_id=0,
                projection_dim=2048,
                text_config={
                    "hidden_activation": "gelu_pytorch_tanh",
                    "hidden_size": 2048,
                    "intermediate_size": 16384,
                    "model_type": "gemma",
                    "num_attention_heads": 8,
                    "num_hidden_layers": 18,
                    "num_image_tokens": 256,
                    "num_key_value_heads": 1,
                    "torch_dtype": "float32",
                    "vocab_size": 257152,
                    "use_adarms": False,
                    "adarms_cond_dim": None,
                },
                vision_config={
                    "hidden_size": 1152,
                    "intermediate_size": 4304,
                    "model_type": "siglip_vision_model",
                    "num_attention_heads": 16,
                    "num_hidden_layers": 27,
                    "num_image_tokens": 256,
                    "patch_size": 14,
                    "projection_dim": 2048,
                    "projector_hidden_act": "gelu_fast",
                    "torch_dtype": "float32",
                    "vision_use_head": False,
                },
            )
        elif isinstance(self.paligemma_config, dict):
            # Override Pi0 default config for PaliGemma
            if "model_type" not in gemma_expert_config:
                paligemma_config["model_type"] = "paligemma"

            cfg_cls = CONFIG_MAPPING[paligemma_config["model_type"]]
            self.paligemma_config = cfg_cls(**paligemma_config)

        if gemma_expert_config is None:
            # Default config from Pi0
            self.gemma_expert_config = CONFIG_MAPPING["gemma"](
                attention_bias=False,
                attention_dropout=0.0,
                bos_token_id=2,
                eos_token_id=1,
                head_dim=256,
                hidden_act="gelu_pytorch_tanh",
                hidden_activation="gelu_pytorch_tanh",
                hidden_size=1024,
                initializer_range=0.02,
                intermediate_size=4096,
                max_position_embeddings=8192,
                model_type="gemma",
                num_attention_heads=8,
                num_hidden_layers=18,
                num_key_value_heads=1,
                pad_token_id=0,
                rms_norm_eps=1e-06,
                rope_theta=10000.0,
                torch_dtype="float32",
                use_adarms=True,
                adarms_cond_dim=1024,
                transformers_version="4.48.1",
                use_cache=True,
                vocab_size=257152,
            )
        elif isinstance(self.gemma_expert_config, dict):
            # Override Pi0 default config for Gemma Expert
            if "model_type" not in gemma_expert_config:
                gemma_expert_config["model_type"] = "gemma"

            cfg_cls = CONFIG_MAPPING[paligemma_config["model_type"]]
            self.gemma_expert_config = cfg_cls(**gemma_expert_config)

        super().__init__(**kwargs)

    def __post_init__(self):
        """Validates configuration parameters."""
        super().__post_init__()
        if self.train_expert_only and not self.freeze_vision_encoder:
            raise ValueError(
                "You set `freeze_vision_encoder=False` and `train_expert_only=True` which are not compatible."
            )

        if self.attention_implementation not in ["eager", "fa2"]:
            raise ValueError(
                f"Wrong value provided for `attention_implementation` ({self.attention_implementation}). Expected 'eager' or 'fa2'."
            )


class PaliGemmaWithExpertModel(PreTrainedModel):
    """PaliGemma model with an additional expert module for action generation."""

    config_class = PaliGemmaWithExpertConfig

    def __init__(self, config: PaliGemmaWithExpertConfig):
        """Initializes the PaliGemmaWithExpertModel.

        Args:
            config: Configuration object for the model.
        """
        super().__init__(config=config)
        self.config = config

        if config.load_pretrained_paligemma:
            self.paligemma = PaliGemmaForConditionalGeneration.from_pretrained("google/paligemma-3b-pt-224")
        else:
            self.paligemma = PaliGemmaForConditionalGeneration(config=config.paligemma_config)
        self.gemma_expert = GemmaForCausalLM(config=config.gemma_expert_config)
        # Remove unused embed_tokens
        self.gemma_expert.model.embed_tokens = None

        # Learned embedding layer for discrete actions
        # Embedding dimension matches expert model hidden size
        self.discrete_action_embedding = nn.Embedding(
            num_embeddings=config.discrete_action_vocab_size,
            embedding_dim=config.paligemma_config.text_config.hidden_size,
            padding_idx=0,  # 0 is used for padding in pad_fast_tokens
        )

        # discrete action head that maps to action vocab size and not language vocab size
        self.da_head = nn.Linear(
            in_features=config.paligemma_config.text_config.hidden_size,
            out_features=config.discrete_action_vocab_size,
        )

        self.dropout = nn.Dropout(config.dropout)

        if not torch.compiler.is_compiling():  # Only cast to bfloat16 if not compiling
            self.to_bfloat16_like_physical_intelligence()
        self.set_requires_grad()

    def set_requires_grad(self) -> None:
        """Sets the requires_grad attribute for model parameters based on configuration."""
        if self.config.freeze_vision_encoder:
            self.paligemma.vision_tower.eval()
            for params in self.paligemma.vision_tower.parameters():
                params.requires_grad = False

        if self.config.train_expert_only:
            self.paligemma.eval()
            for params in self.paligemma.parameters():
                params.requires_grad = False
            for param in self.da_head.parameters():
                param.requires_grad = False
            for param in self.discrete_action_embedding.parameters():
                param.requires_grad = False

    def train(self, mode: bool = True) -> None:
        """Sets the module in training mode.

        Args:
            mode: whether to set training mode (True) or evaluation mode (False). Defaults to True.
        """
        super().train(mode)

        if self.config.freeze_vision_encoder:
            self.paligemma.vision_tower.eval()

        if self.config.train_expert_only:
            self.paligemma.eval()

    def to_bfloat16_like_physical_intelligence(self) -> None:
        """Casts specific model components to bfloat16 dtype."""
        self.paligemma = self.paligemma.to(dtype=torch.bfloat16)

        params_to_change_dtype = [
            "language_model.model.layers",
            "gemma_expert.model.layers",
            "vision_tower",
            "multi_modal",
        ]
        for name, param in self.named_parameters():
            if any(selector in name for selector in params_to_change_dtype):
                param.data = param.data.to(dtype=torch.bfloat16)

    def embed_image(self, image: torch.Tensor) -> torch.Tensor:
        """Computes image embeddings.

        Args:
            image: Input image tensor.

        Returns:
            torch.Tensor: Image embeddings.
        """
        # Handle different transformers versions
        if hasattr(self.paligemma, "get_image_features"):
            return self.paligemma.get_image_features(image)
        else:
            return self.paligemma.model.get_image_features(image)

    def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor:
        """Embeds language tokens.

        Args:
            tokens: Input token indices.

        Returns:
            torch.Tensor: Token embeddings.
        """
        return self.paligemma.language_model.embed_tokens(tokens)

    def embed_discrete_actions(self, actions: torch.Tensor) -> torch.Tensor:
        """Embeds discrete action tokens.

        Args:
            actions: Input discrete action indices.

        Returns:
            torch.Tensor: Action embeddings.
        """
        # Ensure actions are long integers for embedding lookup
        if actions.dtype != torch.long:
            actions = actions.long()

        # Apply embedding layer
        embedded = self.discrete_action_embedding(actions)

        return embedded

    # TODO: break down this huge forward into modules or functions
    def forward(
        self,
        attention_mask: torch.Tensor | None = None,
        position_ids: torch.LongTensor | None = None,
        past_key_values: list[torch.FloatTensor] | Cache | None = None,
        inputs_embeds: list[torch.FloatTensor] = None,
        n_cross_att_tokens: int | None = None,
        use_cache: bool | None = None,
        fill_kv_cache: bool | None = None,
        adarms_cond: list[torch.Tensor] | None = None,
    ) -> tuple[list[torch.FloatTensor | None], list[torch.FloatTensor] | Cache | None]:
        """Forward pass of the model.

        Args:
            attention_mask: Attention mask tensor.
            position_ids: Position IDs tensor.
            past_key_values: Past key values for caching.
            inputs_embeds: List of input embeddings for the different model parts.
            n_cross_att_tokens: Number of cross-attention tokens.
            use_cache: Whether to use KV cache.
            fill_kv_cache: Whether to fill the KV cache.
            adarms_cond: List of AdaRMS conditioning tensors.

        Returns:
            tuple: A tuple containing:
                - outputs_embeds: List of output embeddings.
                - past_key_values: Updated past key values.

        Raises:
            ValueError: If `n_cross_att_tokens` is not provided when `fill_kv_cache` is True.
        """
        if adarms_cond is None:
            adarms_cond = [None, None]

        models = [self.paligemma.language_model, self.gemma_expert.model]

        for hidden_states in inputs_embeds:
            # TODO this is very inefficient
            # dtype is always the same, batch size too (if > 1 len)
            # device could be trickier in multi gpu edge cases but that's it
            if hidden_states is None:
                continue
            batch_size = hidden_states.shape[0]

        # RMSNorm
        num_layers = self.paligemma.config.text_config.num_hidden_layers
        head_dim = self.paligemma.config.text_config.head_dim
        for layer_idx in range(num_layers):
            query_states = []
            key_states = []
            value_states = []
            gates = []
            for i, hidden_states in enumerate(inputs_embeds):
                if hidden_states is None:
                    gates.append(None)
                    continue
                layer = models[i].layers[layer_idx]
                # normalizer = torch.tensor(models[i].config.hidden_size**0.5, dtype=hidden_states.dtype)
                # hidden_states = hidden_states * normalizer
                hidden_states, gate = layer.input_layernorm(hidden_states, cond=adarms_cond[i])
                gates.append(gate)
                input_shape = hidden_states.shape[:-1]
                hidden_shape = (*input_shape, -1, layer.self_attn.head_dim)

                hidden_states = hidden_states.to(dtype=_preferred_dtype())
                query_state = layer.self_attn.q_proj(hidden_states).view(hidden_shape)
                key_state = layer.self_attn.k_proj(hidden_states).view(hidden_shape)
                value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape)

                query_states.append(query_state)
                key_states.append(key_state)
                value_states.append(value_state)

            # B,L,H,D with L sequence length, H number of heads, D head dim
            # concatenate on the number of embeddings/tokens
            query_states = torch.cat(query_states, dim=1)
            key_states = torch.cat(key_states, dim=1)
            value_states = torch.cat(value_states, dim=1)

            query_states = apply_rope(query_states, position_ids)
            key_states = apply_rope(key_states, position_ids)

            if use_cache:
                # TODO here, some optimization can be done - similar to a `StaticCache` we can declare the `max_len` before.
                # so we create an empty cache, with just one cuda malloc, and if (in autoregressive case) we reach
                # the max len, then we (for instance) double the cache size. This implementation already exists
                # in `transformers`. (molbap)
                key_states = torch.cat([past_key_values[layer_idx]["key_states"], key_states], dim=1)
                value_states = torch.cat([past_key_values[layer_idx]["value_states"], value_states], dim=1)
            if fill_kv_cache:
                if past_key_values is None:
                    past_key_values = {}
                if n_cross_att_tokens is None:
                    raise ValueError("n_cross_att_tokens must be provided when fill_kv_cache is True")
                past_key_values[layer_idx] = {
                    # save the first n_cross_att_tokens for action expert cross attention
                    "key_states": key_states[:, :n_cross_att_tokens, :, :],
                    "value_states": value_states[:, :n_cross_att_tokens, :, :],
                }

            attention_interface = self.get_attention_interface()
            att_output = attention_interface(
                attention_mask, batch_size, head_dim, query_states, key_states, value_states
            )
            att_output = att_output.to(dtype=_preferred_dtype())

            # first part of att_output is prefix (up to sequence length, [:, 0:prefix_seq_len])
            outputs_embeds = []
            start = 0
            for i, hidden_states in enumerate(inputs_embeds):
                layer = models[i].layers[layer_idx]

                if hidden_states is not None:
                    end = start + hidden_states.shape[1]

                    if att_output.dtype != layer.self_attn.o_proj.weight.dtype:
                        att_output = att_output.to(layer.self_attn.o_proj.weight.dtype)
                    out_emb = layer.self_attn.o_proj(att_output[:, start:end])

                    out_emb = self.dropout(out_emb)

                    # first residual
                    out_emb = modeling_gemma._gated_residual(hidden_states, out_emb, gates[i])  # noqa: SLF001
                    after_first_residual = out_emb.clone()

                    out_emb, gate = layer.post_attention_layernorm(out_emb, cond=adarms_cond[i])
                    out_emb = layer.mlp(out_emb)

                    out_emb = self.dropout(out_emb)

                    # second residual
                    out_emb = modeling_gemma._gated_residual(after_first_residual, out_emb, gate)  # noqa: SLF001

                    outputs_embeds.append(out_emb)

                    start = end
                else:
                    outputs_embeds.append(None)

            inputs_embeds = outputs_embeds

        # final norm
        outputs_embeds = []
        for i, hidden_states in enumerate(inputs_embeds):
            if hidden_states is not None:
                out_emb, _ = models[i].norm(hidden_states, cond=adarms_cond[i])
                outputs_embeds.append(out_emb)
            else:
                outputs_embeds.append(None)

        return outputs_embeds, past_key_values

    def get_attention_interface(self):
        """Returns the attention implementation function based on config.

        Returns:
            callable: The attention function to use.
        """
        return self.eager_attention_forward

    def eager_attention_forward(
        self,
        attention_mask: torch.Tensor,
        batch_size: int,
        head_dim: int,
        query_states: torch.Tensor,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
    ) -> torch.Tensor:
        """Eager attention forward pass using standard matrix multiplications.

        Args:
            attention_mask: Attention mask tensor.
            batch_size: Batch size.
            head_dim: Head dimension.
            query_states: Query states tensor.
            key_states: Key states tensor.
            value_states: Value states tensor.

        Returns:
            torch.Tensor: Attention output.
        """
        num_att_heads = self.config.paligemma_config.text_config.num_attention_heads
        num_key_value_heads = self.config.paligemma_config.text_config.num_key_value_heads
        num_key_value_groups = num_att_heads // num_key_value_heads

        # query_states: batch_size, sequence_length, num_att_head, head_dim
        # key_states: batch_size, sequence_length, num_key_value_head, head_dim
        # value_states: batch_size, sequence_length, num_key_value_head, head_dim
        sequence_length = key_states.shape[1]

        key_states = key_states[:, :, :, None, :].expand(
            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
        )
        key_states = key_states.reshape(
            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
        )

        value_states = value_states[:, :, :, None, :].expand(
            batch_size, sequence_length, num_key_value_heads, num_key_value_groups, head_dim
        )
        value_states = value_states.reshape(
            batch_size, sequence_length, num_key_value_heads * num_key_value_groups, head_dim
        )

        # Attention here is upcasted to float32 to match the original eager implementation.

        query_states = query_states.to(dtype=torch.float32)
        key_states = key_states.to(dtype=torch.float32)

        query_states = query_states.transpose(1, 2)
        key_states = key_states.transpose(1, 2)

        att_weights = torch.matmul(query_states, key_states.transpose(2, 3))
        att_weights *= head_dim**-0.5
        big_neg = -2.3819763e38  # See gemma/modules.py

        masked_att_weights = torch.where(attention_mask[:, None, :, :], att_weights, big_neg)

        probs = nn.functional.softmax(masked_att_weights, dim=-1)
        probs = probs.to(dtype=value_states.dtype)

        # probs: batch_size, num_key_value_head, num_att_head, sequence_length, sequence_length
        # value_states: batch_size, sequence_length, num_att_heads, head_dim

        att_output = torch.matmul(probs, value_states.permute(0, 2, 1, 3))

        att_output = att_output.permute(0, 2, 1, 3)
        # we use -1 because sequence length can change
        att_output = att_output.reshape(batch_size, -1, num_key_value_heads * num_key_value_groups * head_dim)

        return att_output
