import torch
import math
import numpy as np


def get_1d_sincos_pos_emb_from_grid(embed_dim, pos, device="cpu"):
    """
    Generate 1D sinusoidal positional embeddings from grid positions.
    
    Args:
        embed_dim (int): The embedding dimension (must be even).
        pos (torch.Tensor): The grid positions (e.g., [0, 1, 2, ..., v-1]).
                           Shape: [b * gh * gw] or [batch_size, sequence_length].
        device (str): Device for the output tensor.
    
    Returns:
        torch.Tensor: Sinusoidal positional embeddings.
                      Shape: [len(pos), embed_dim]
    """
    assert embed_dim % 2 == 0, "Embedding dimension must be even for sine and cosine."
    
    # Convert positions to float
    pos = pos.float()
    
    # Compute the sinusoidal frequencies
    dim = torch.arange(embed_dim // 2, dtype=torch.float32, device=device)  # [0, 1, ..., embed_dim // 2 - 1]
    freq = 1.0 / (10000 ** (dim / (embed_dim // 2)))  # Scale frequencies logarithmically
    
    # Calculate sine and cosine embeddings
    pos_emb_sin = torch.sin(pos[:, None] * freq)  # Shape: [len(pos), embed_dim // 2]
    pos_emb_cos = torch.cos(pos[:, None] * freq)  # Shape: [len(pos), embed_dim // 2]
    
    # Concatenate sine and cosine along the last dimension
    pos_emb = torch.cat([pos_emb_sin, pos_emb_cos], dim=-1)  # Shape: [len(pos), embed_dim]
    
    return pos_emb


def get_2d_sincos_pos_embed(embed_dim, grid_size, device="cpu"):
    """
    Generate 2D sine-cosine positional embeddings with separate grid height and width.
    
    Args:
        embed_dim (int): The embedding dimension.
        grid_size (tuple): Tuple specifying the grid height and width (grid_h, grid_w).
        cls_token (bool): Whether to include a [CLS] token embedding.
        device (str): The device to place the embeddings on.
    
    Returns:
        torch.Tensor: Positional embeddings of shape [grid_h*grid_w, embed_dim]
    """
    grid_h, grid_w = grid_size  # Unpack grid dimensions

    # Create the grid for height and width
    grid_h = torch.arange(grid_h, dtype=torch.float32, device=device)
    grid_w = torch.arange(grid_w, dtype=torch.float32, device=device)
    grid = torch.meshgrid(grid_w, grid_h, indexing="ij")  # w goes first
    grid = torch.stack(grid, dim=0)  # Shape: [2, grid_h, grid_w]

    # Reshape grid to [2, 1, grid_h, grid_w]
    grid = grid.view(2, 1, grid.size(1), grid.size(2))

    # Get the positional embeddings from the grid
    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid, device=device)  # Shape: [grid_h*grid_w, embed_dim]

    return pos_embed


def get_2d_sincos_pos_embed_from_grid(embed_dim, grid, device="cpu"):
    """
    Generate 2D sine-cosine positional embeddings from a grid with separate height and width.
    
    Args:
        embed_dim (int): The embedding dimension.
        grid (torch.Tensor): The grid of shape [2, 1, grid_h, grid_w].
        device (str): The device to place the embeddings on.
    
    Returns:
        torch.Tensor: Positional embeddings of shape [grid_h*grid_w, embed_dim].
    """
    assert embed_dim % 2 == 0, "Embedding dimension must be even."

    # Use half of the dimensions for each grid dimension
    grid_h = grid[0].view(-1)  # Flatten grid height dimension: [H*W]
    grid_w = grid[1].view(-1)  # Flatten grid width dimension: [H*W]

    # Generate 1D sine-cosine embeddings for grid_h and grid_w
    emb_h = get_1d_sincos_pos_emb_from_grid(embed_dim // 2, grid_h, device=device)  # Shape: [H*W, D/2]
    emb_w = get_1d_sincos_pos_emb_from_grid(embed_dim // 2, grid_w, device=device)  # Shape: [H*W, D/2]

    # Concatenate along the last dimension
    pos_embed = torch.cat([emb_h, emb_w], dim=-1)  # Shape: [H*W, D]
    return pos_embed


def rope(positions: torch.Tensor, d: int, device='cpu') -> torch.Tensor:
    """
    Given a batch of positions in [0,1], compute RoPE-style
    sine-cosine embeddings in dimension d (must be even).
    
    positions: (B, N) tensor of float positions in [0,1].
    d: int, dimension of the embedding (should be even).
    Returns:
      embeddings: (B, N, d) tensor of float embeddings.
    """
    # positions shape: [B, N]
    B, N = positions.shape
    half_d = d // 2

    # Expand positions to shape [B, N, 1]
    positions_3d = positions.unsqueeze(-1)  # [B, N, 1]

    # Prepare index and frequency tensors
    # idx => [1, 1, half_d]
    idx = torch.arange(half_d, device=device).view(1, 1, -1)
    # freqs => [1, 1, half_d], broadcast to [B, N, half_d]
    freqs = torch.pow(10000.0, -2.0 * idx / d)

    # angle => [B, N, half_d]
    angle = positions_3d.to(device) * freqs

    # Compute sine and cosine => each [B, N, half_d]
    sin_part = angle.sin()
    cos_part = angle.cos()

    # Interleave sine and cosine along the last dimension => [B, N, d]
    embeddings = torch.empty(B, N, d, device=device, dtype=positions.dtype)
    embeddings[..., 0::2] = sin_part
    embeddings[..., 1::2] = cos_part

    return embeddings