# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates and/or its affiliates
#
# 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.

import numpy as np
import torch
from diffusers.models.embeddings import get_1d_rotary_pos_embed as _get_1d_rotary_pos_embed
from einops import repeat
from torch import nn

from vllm_omni.platforms import current_omni_platform


def apply_real_rotary_emb(x: torch.Tensor, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor) -> torch.Tensor:
    """
    Apply real-valued rotary embeddings to input tensor.

    Args:
        x: Input tensor of shape [..., seq_len, num_heads, dim] or [..., seq_len, dim]
        freqs_cos: Cosine frequencies of shape [batch, seq_len, dim] or [seq_len, dim]
        freqs_sin: Sine frequencies of shape [batch, seq_len, dim] or [seq_len, dim]

    Returns:
        Tensor with rotary embeddings applied
    """
    # x: [batch, seq_len, num_heads, dim] or [batch, seq_len, dim]
    # freqs_cos: [batch, seq_len, dim] or [seq_len, dim]
    # freqs_sin: [batch, seq_len, dim] or [seq_len, dim]

    x_shape = x.shape
    if len(x_shape) == 4:
        batch, seq_len, num_heads, dim = x_shape
        x_reshaped = x.view(batch, seq_len, num_heads, dim // 2, 2)
    elif len(x_shape) == 3:
        batch, seq_len, dim = x_shape
        num_heads = None
        x_reshaped = x.view(batch, seq_len, dim // 2, 2)
    else:
        raise ValueError(f"Unsupported x shape: {x.shape}")

    # freqs_cos/sin: [batch, seq_len, dim] or [seq_len, dim]
    # Expand freqs_cos/sin to [batch, seq_len, dim] if needed
    if freqs_cos.dim() == 2:
        # [seq_len, dim] -> [1, seq_len, dim]
        freqs_cos = freqs_cos.unsqueeze(0)
        freqs_sin = freqs_sin.unsqueeze(0)
    if freqs_cos.shape[0] == 1 and batch > 1:
        freqs_cos = freqs_cos.expand(batch, -1, -1)
        freqs_sin = freqs_sin.expand(batch, -1, -1)

    # Reshape freqs to [batch, seq_len, dim//2, 2]
    freqs_cos_reshaped = freqs_cos.view(batch, seq_len, dim // 2, 2)
    freqs_sin_reshaped = freqs_sin.view(batch, seq_len, dim // 2, 2)

    cos_1 = freqs_cos_reshaped[..., 0]  # [batch, seq_len, dim//2]
    cos_2 = freqs_cos_reshaped[..., 1]  # [batch, seq_len, dim//2]
    sin_1 = freqs_sin_reshaped[..., 0]  # [batch, seq_len, dim//2]
    sin_2 = freqs_sin_reshaped[..., 1]  # [batch, seq_len, dim//2]

    # Broadcast cos/sin to match x_reshaped
    if len(x_shape) == 4:
        # [batch, seq_len, 1, dim//2]
        cos_1 = cos_1.unsqueeze(2)
        cos_2 = cos_2.unsqueeze(2)
        sin_1 = sin_1.unsqueeze(2)
        sin_2 = sin_2.unsqueeze(2)

    x1 = x_reshaped[..., 0]  # [..., seq_len, num_heads, dim//2] or [..., seq_len, dim//2]
    x2 = x_reshaped[..., 1]  # same

    out1 = x1 * cos_1 - x2 * sin_1
    out2 = x1 * sin_2 + x2 * cos_2

    out = torch.stack([out1, out2], dim=-1)
    return out.view(*x_shape)


def get_1d_rotary_pos_embed_real(
    dim: int,
    pos: np.ndarray | int,
    theta: float = 10000.0,
    linear_factor: float = 1.0,
    ntk_factor: float = 1.0,
    freqs_dtype: torch.dtype = torch.float32,
):
    freqs_cos, freqs_sin = _get_1d_rotary_pos_embed(
        dim,
        pos,
        theta=theta,
        use_real=True,
        linear_factor=linear_factor,
        ntk_factor=ntk_factor,
        repeat_interleave_real=True,
        freqs_dtype=freqs_dtype,
    )
    return freqs_cos, freqs_sin


class RotaryPosEmbedReal(nn.Module):
    def __init__(
        self, theta: int, axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], patch_size: int = 2
    ):
        super().__init__()
        self.theta = theta
        self.axes_dim = axes_dim
        self.axes_lens = axes_lens
        self.patch_size = patch_size

    @staticmethod
    def get_freqs_real(
        axes_dim: tuple[int, int, int], axes_lens: tuple[int, int, int], theta: int
    ) -> list[tuple[torch.Tensor, torch.Tensor]]:
        freqs_real = []
        freqs_dtype = torch.float64 if current_omni_platform.supports_float64() else torch.float32
        for i, (d, e) in enumerate(zip(axes_dim, axes_lens)):
            cos_emb, sin_emb = get_1d_rotary_pos_embed_real(d, e, theta=theta, freqs_dtype=freqs_dtype)
            freqs_real.append((cos_emb, sin_emb))
        return freqs_real

    def _get_freqs_real(
        self, freqs_real: list[tuple[torch.Tensor, torch.Tensor]], ids: torch.Tensor
    ) -> tuple[torch.Tensor, torch.Tensor]:
        device = ids.device
        if ids.device.type == "mps":
            ids = ids.to("cpu")

        cos_result = []
        sin_result = []
        for i in range(len(self.axes_dim)):
            freqs_cos, freqs_sin = freqs_real[i]
            freqs_cos = freqs_cos.to(ids.device)
            freqs_sin = freqs_sin.to(ids.device)
            index = ids[:, :, i : i + 1].repeat(1, 1, freqs_cos.shape[-1]).to(torch.int64)
            cos_result.append(torch.gather(freqs_cos.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index))
            sin_result.append(torch.gather(freqs_sin.unsqueeze(0).repeat(index.shape[0], 1, 1), dim=1, index=index))

        combined_cos = torch.cat(cos_result, dim=-1).to(device)
        combined_sin = torch.cat(sin_result, dim=-1).to(device)
        return combined_cos, combined_sin

    def forward(
        self, freqs_real, attention_mask, l_effective_ref_img_len, l_effective_img_len, ref_img_sizes, img_sizes, device
    ):
        batch_size = len(attention_mask)
        p = self.patch_size

        encoder_seq_len = attention_mask.shape[1]
        l_effective_cap_len = attention_mask.sum(dim=1).tolist()
        l_effective_cap_len = [int(len) for len in l_effective_cap_len]
        seq_lengths = [
            int(cap_len + sum(ref_img_len) + img_len)
            for cap_len, ref_img_len, img_len in zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len)
        ]

        max_seq_len = max(seq_lengths)
        max_ref_img_len = max([int(sum(ref_img_len)) for ref_img_len in l_effective_ref_img_len])
        max_img_len = int(max(l_effective_img_len))

        # Create position IDs
        position_ids = torch.zeros(batch_size, int(max_seq_len), 3, dtype=torch.int32, device=device)

        for i, (cap_seq_len, seq_len) in enumerate(zip(l_effective_cap_len, seq_lengths)):
            # add text position ids
            position_ids[i, :cap_seq_len] = repeat(
                torch.arange(cap_seq_len, dtype=torch.int32, device=device), "l -> l 3"
            )
            pe_shift = cap_seq_len
            pe_shift_len = cap_seq_len

            if ref_img_sizes[i] is not None:
                for ref_img_size, ref_img_len in zip(ref_img_sizes[i], l_effective_ref_img_len[i]):
                    H, W = ref_img_size
                    ref_H_tokens, ref_W_tokens = H // p, W // p
                    assert ref_H_tokens * ref_W_tokens == ref_img_len
                    # add image position ids

                    row_ids = repeat(
                        torch.arange(ref_H_tokens, dtype=torch.int32, device=device), "h -> h w", w=ref_W_tokens
                    ).flatten()
                    col_ids = repeat(
                        torch.arange(ref_W_tokens, dtype=torch.int32, device=device), "w -> h w", h=ref_H_tokens
                    ).flatten()
                    position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 0] = pe_shift
                    position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 1] = row_ids
                    position_ids[i, pe_shift_len : pe_shift_len + ref_img_len, 2] = col_ids

                    pe_shift += max(ref_H_tokens, ref_W_tokens)
                    pe_shift_len += ref_img_len

            H, W = img_sizes[i]
            H_tokens, W_tokens = H // p, W // p
            assert H_tokens * W_tokens == l_effective_img_len[i]

            row_ids = repeat(torch.arange(H_tokens, dtype=torch.int32, device=device), "h -> h w", w=W_tokens).flatten()
            col_ids = repeat(torch.arange(W_tokens, dtype=torch.int32, device=device), "w -> h w", h=H_tokens).flatten()

            assert pe_shift_len + l_effective_img_len[i] == seq_len
            position_ids[i, pe_shift_len:seq_len, 0] = pe_shift
            position_ids[i, pe_shift_len:seq_len, 1] = row_ids
            position_ids[i, pe_shift_len:seq_len, 2] = col_ids

        # Get combined rotary embeddings (real version)
        freqs_cos, freqs_sin = self._get_freqs_real(freqs_real, position_ids)

        # create separate rotary embeddings for captions and images
        cap_freqs_cos = torch.zeros(
            batch_size, encoder_seq_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype
        )
        cap_freqs_sin = torch.zeros(
            batch_size, encoder_seq_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype
        )
        ref_img_freqs_cos = torch.zeros(
            batch_size, max_ref_img_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype
        )
        ref_img_freqs_sin = torch.zeros(
            batch_size, max_ref_img_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype
        )
        img_freqs_cos = torch.zeros(batch_size, max_img_len, freqs_cos.shape[-1], device=device, dtype=freqs_cos.dtype)
        img_freqs_sin = torch.zeros(batch_size, max_img_len, freqs_sin.shape[-1], device=device, dtype=freqs_sin.dtype)

        for i, (cap_seq_len, ref_img_len, img_len, seq_len) in enumerate(
            zip(l_effective_cap_len, l_effective_ref_img_len, l_effective_img_len, seq_lengths)
        ):
            cap_freqs_cos[i, :cap_seq_len] = freqs_cos[i, :cap_seq_len]
            cap_freqs_sin[i, :cap_seq_len] = freqs_sin[i, :cap_seq_len]
            ref_img_freqs_cos[i, : sum(ref_img_len)] = freqs_cos[i, cap_seq_len : cap_seq_len + sum(ref_img_len)]
            ref_img_freqs_sin[i, : sum(ref_img_len)] = freqs_sin[i, cap_seq_len : cap_seq_len + sum(ref_img_len)]
            img_freqs_cos[i, :img_len] = freqs_cos[
                i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len
            ]
            img_freqs_sin[i, :img_len] = freqs_sin[
                i, cap_seq_len + sum(ref_img_len) : cap_seq_len + sum(ref_img_len) + img_len
            ]

        return (
            (cap_freqs_cos, cap_freqs_sin),
            (ref_img_freqs_cos, ref_img_freqs_sin),
            (img_freqs_cos, img_freqs_sin),
            (freqs_cos, freqs_sin),
            l_effective_cap_len,
            seq_lengths,
        )