# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# Position embedding utils
# --------------------------------------------------------

import numpy as np
import torch



# --------------------------------------------------------
# 2D sine-cosine position embedding
# References:
# Transformer: https://github.com/tensorflow/models/blob/master/official/nlp/transformer/model_utils.py
# MoCo v3: https://github.com/facebookresearch/moco-v3
# --------------------------------------------------------
def get_1d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
    """
    grid_size: int of the grid height and width
    return:
    pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
    """
    # grid = grid.reshape([2, 1, grid_size, grid_size])
    grid = np.arange(grid_size, dtype=np.float32)
    grid = grid.reshape([1, 1, grid_size])
    pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid)
    if cls_token:
        pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
    return pos_embed

def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
    """
    embed_dim: output dimension for each position
    pos: a list of positions to be encoded: size (M,)
    out: (M, D)
    """
    assert embed_dim % 2 == 0
    omega = np.arange(embed_dim // 2, dtype=np.float32)
    omega /= embed_dim / 2.
    omega = 1. / 10000**omega  # (D/2,)

    pos = pos.reshape(-1)  # (M,)
    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product

    emb_sin = np.sin(out) # (M, D/2)
    emb_cos = np.cos(out) # (M, D/2)

    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
    return emb



def flip_padded_sequences(batch, mask):
    """
    Flips the valid part of each sequence in a padded batch, keeping pads at the end.
    
    Args:
        batch (Tensor): Input tensor of size (B, L).
        lengths (Tensor, optional): Lengths of size (B,). If None, calc from mask.
        padding_value (int): Value used for padding.
        
    Returns:
        Tensor: The flipped batch.
    """
    B, L, D = batch.shape
    device = batch.device
    
    # 1. Get Sequence Lengths
    lengths = mask.sum(dim=1).long()
    
    # 2. Generate Index Grid
    # Shape: (L,) -> [0, 1, 2, ..., L-1]
    range_tensor = torch.arange(L, device=device)
    
    # Expand to (B, L)
    # [[0, 1, ...], [0, 1, ...], ...]
    idx_grid = range_tensor.unsqueeze(0).expand(B, L)
    
    # 3. Calculate "Inverse" Indices for valid positions
    # If len=3, indices [0, 1, 2] should become [2, 1, 0]
    # Formula: flipped_index = (Length - 1) - current_index
    lens_expanded = lengths.unsqueeze(1)
    flipped_indices = (lens_expanded - 1) - idx_grid
    
    # 4. Create the final Gather Indices
    # We only want to use the flipped indices where the data is valid.
    # Where it is padding, we keep the original index (so we just gather the 0s back).
    mask = idx_grid < lens_expanded
    
    # Use torch.where: if valid, use flipped_indices, else use original idx_grid
    gather_indices = torch.where(mask, flipped_indices, idx_grid)
    
    # 5. Gather
    result = torch.gather(batch, 1, gather_indices.unsqueeze(-1).repeat(1, 1, D))
    
    return result