# Copyright (c) 2022 Alpha-VL
# --------------------------------------------------------
# References:
# timm: https://github.com/rwightman/pytorch-image-models/tree/master/timm
# DeiT: https://github.com/facebookresearch/deit
# --------------------------------------------------------
import torch
import torch.nn as nn


from timm.models.layers import DropPath

from flash_attn import flash_attn_func
from einops import rearrange, repeat
import math

import torch.nn.functional as F

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.1):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

class GatedMLP(nn.Module):
    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=F.silu,        
        multiple_of=128, 
        drop=0.3       
    ):
        
        super().__init__()
        out_features = out_features if out_features is not None else in_features
        hidden_features = (
            hidden_features if hidden_features is not None else int(8 * in_features / 3)
        )
        hidden_features = (hidden_features + multiple_of - 1) // multiple_of * multiple_of
        self.fc1 = nn.Linear(in_features, 2 * hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        y = self.fc1(x)
        y, gate = y.chunk(2, dim=-1)
        y = y * self.act(gate)
        y = self.drop(y)
        y = self.fc2(y)
        y = self.drop(y)
        return y


class Attention(nn.Module):
    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
        self.scale = qk_scale or head_dim ** -0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x, mask=None):
        B, N, C = x.shape
        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]   # make torchscript happy (cannot use tensor as tuple)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        
        # Apply mask if provided
        if mask is not None:
            # mask shape: (B, N) or (B, 1, 1, N) for broadcasting
            if mask.dim() == 2:
                mask = mask.unsqueeze(1).unsqueeze(1)  # (B, 1, 1, N)
            # Convert boolean mask to attention mask: False->0 (keep), True->-inf (mask out)
            attn_mask = torch.where(mask, torch.tensor(float('-inf')), torch.tensor(0.0))
            attn = attn + attn_mask
        
        attn = attn.softmax(dim=-1)
        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x


class Block(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0.1, attn_drop=0.1,
                 drop_path=0.1, act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        # self.mlp = GatedMLP(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, x, mask=None):
        x = x + self.drop_path(self.attn(self.norm1(x), mask=mask))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x


class MoEBlock(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0.1, attn_drop=0.1,
                 drop_path=0.1, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
                 seq_len=375, num_experts=3, num_experts_per=2):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)

        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        
        self.num_experts = num_experts
        self.num_experts_per = num_experts_per
        self.gate = nn.Linear(
            dim * seq_len, num_experts, bias=False)
        # Patchwise gating, worse performance.
        # self.gate = nn.Linear(
        #     dim, num_experts, bias=False)
        self.experts = nn.ModuleList(
            [Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) 
            for _ in range(num_experts)])
        

    def forward(self, x, mask=None):
        x = x + self.drop_path(self.attn(self.norm1(x), mask=mask))
        x_orig = x
        x = self.norm2(x)        
        orig_shape = x.shape
        x = x.view(x.shape[0], -1) #(B, L*C)
        scores = self.gate(x) #(B, num_experts)
        expert_weights, expert_indices = torch.topk(
            scores, self.num_experts_per, dim=-1) #(B, #num_experts_per)        

        # Weights are normlaized
        expert_weights = expert_weights.softmax(dim=-1) #(B, #num_experts_per)
        expert_weights = expert_weights.unsqueeze(1).repeat(1, orig_shape[1], 1) #(B, L, #num_experts_per)
        expert_indices = expert_indices.unsqueeze(1).repeat(1, orig_shape[1], 1) #(B, L, #num_experts_per)        
        
        x = x.view(orig_shape[0], orig_shape[1], 1, -1)  # (B, L, 1, C)
        x = x.repeat_interleave(self.num_experts_per, dim=2) # (B, L, num_experts_per, C)
        y = torch.empty_like(x, dtype=x.dtype) #(B, L, num_experts_per, C)
        for i, expert in enumerate(self.experts):
            y[expert_indices == i] = expert(x[expert_indices == i]).to(x.dtype)            
        
        y = (y.view(*expert_weights.shape, -1) * expert_weights.unsqueeze(-1)).sum(dim=2)
        x = x_orig + self.drop_path(y.view(*orig_shape))

        return x, scores


class BlockFlashAttn(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0.1, attn_drop=0.1,
                 drop_path=0.1, act_layer=nn.GELU, norm_layer=nn.LayerNorm):
        super().__init__()
        self.norm1_q = norm_layer(dim)
        self.norm1_kv = norm_layer(dim)
        # self.attn = Attention(
        #     dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
        self.q_proj = nn.Linear(dim, dim, bias=qkv_bias)
        self.k_proj = nn.Linear(dim, dim, bias=qkv_bias)
        self.v_proj = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = attn_drop
        self.num_heads = num_heads

        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = GatedMLP(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)

    def forward(self, x, context=None, rope=None, rope2=None, mask=None):
        dropout_p = self.attn_drop if self.training else 0.0        
        x_ = x

        if context is None:
            context = x
        # Project queries, keys, values
        q = self.q_proj(self.norm1_q(x))
        k = self.k_proj(self.norm1_kv(context))
        v = self.v_proj(self.norm1_kv(context))
        # Reshape for flash attention: (batch, seqlen, nheads, headdim)
        B, Lq, C = q.shape
        B, Lkv, C = k.shape
        q = q.reshape(B, Lq, self.num_heads, C // self.num_heads)
        k = k.reshape(B, Lkv, self.num_heads, C // self.num_heads)
        v = v.reshape(B, Lkv, self.num_heads, C // self.num_heads)
        if rope is not None:
            if len(rope.shape) == 3:
                rope = rope.unsqueeze(1).to(dtype=q.dtype)  # (B, 1, seq_len, dim)
            q = rearrange(q, 'b l h d -> b h l d')
            q = rotate(q, rope)
            q = rearrange(q, 'b h l d -> b l h d')    
            
            if rope2 is not None:
                if len(rope2.shape) == 3:
                    rope2 = rope2.unsqueeze(1).to(dtype=q.dtype)  # (B, 1, seq_len, dim)
            else:
                rope2 = rope
                
            k = rearrange(k, 'b l h d -> b h l d')
            k = rotate(k, rope2)
            k = rearrange(k, 'b h l d -> b l h d')
        
        # Stack for flash_attn_func: (q, k, v)
        # flash_attn_func expects (q, k, v) separately
        out = flash_attn_func(q, k, v, dropout_p=dropout_p, causal=False)
        
        out = out.reshape(B, Lq, -1)
        x = x_ + self.drop_path(out)        
        x = x + self.drop_path(self.mlp(self.norm2(x)))        
        return x, None  # For consistency
  
class MoEBlockFlashAttn(nn.Module):

    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0.1, attn_drop=0.1,
                 drop_path=0.1, act_layer=nn.GELU, norm_layer=nn.LayerNorm,
                 seq_len=375, num_experts=3, num_experts_per=2):
        super().__init__()
        self.norm1 = norm_layer(dim)
        # self.attn = Attention(
        #     dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)                
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = attn_drop
        self.num_heads = num_heads

        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        
        self.num_experts = num_experts
        self.num_experts_per = num_experts_per
        self.gate = nn.Linear(
            dim, num_experts, bias=False)
        # Patchwise gating, worse performance.
        # self.gate = nn.Linear(
        #     dim, num_experts, bias=False)
        self.experts = nn.ModuleList(
            [GatedMLP(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) 
            for _ in range(num_experts)])
        

    def forward(self, x, rope=None, mask=None):

        # x = x + self.drop_path(self.attn(self.norm1(x)))
        dropout_p = self.attn_drop if self.training else 0.0
        x_ = x        

        x = self.qkv(self.norm1(x))
        
        # Flash attention requires this shape: qkv: (batch_size, seqlen, 3, nheads, headdim)
        x = x.reshape(x.shape[0], x.shape[1], 3, x.shape[-1] // 3)
        x = x.reshape(x.shape[0], x.shape[1], 3, self.num_heads, x.shape[-1] // self.num_heads)

        q = x[:,:,0]  # (batch_size, seqlen, nheads, headdim)
        k = x[:,:,1]  # (batch_size, seqlen, nheads , headdim)
        v = x[:,:,2]  # (batch_size, seqlen, nheads , headdim)

        if rope is not None:
            if len(rope.shape) == 3:
                rope = rope.unsqueeze(1).to(q.dtype)  # (1, seq_len, 1, dim)
            q = rearrange(q, 'b l h d -> b h l d')
            q = rotate(q, rope)
            q = rearrange(q, 'b h l d -> b l h d')    
            k = rearrange(k, 'b l h d -> b h l d')
            k = rotate(k, rope)
            k = rearrange(k, 'b h l d -> b l h d')
           
        # x, lse, S_dmask = flash_attn_func(q, k, v, dropout_p=dropout_p, causal=False, return_attn_probs=True)
        x= flash_attn_func(q, k, v, dropout_p=dropout_p, causal=False)

        x = x.reshape(x.shape[0], x.shape[1], -1)  # (B, L, C)    

        x = x_ + self.drop_path(x)
        x_orig = x
        x = self.norm2(x)        
        orig_shape = x.shape
        # x = x.view(x.shape[0], -1) #(B, L*C)
        # scores = torch.mean(self.gate(x), dim=1) #(B, num_experts)
        scores = self.gate(x) #(B, L, num_experts)
        # scores = self.gate(x[:,:512*60]) #(B, num_experts)
        expert_weights, expert_indices = torch.topk(
            scores, self.num_experts_per, dim=-1) #(B, L, num_experts_per)        

        # Weights are normlaized
        expert_weights = expert_weights.softmax(dim=-1) #(B, L, num_experts_per)
        # expert_weights = expert_weights.repeat(1, orig_shape[1], 1) #(B, L, #num_experts_per)
        # expert_indices = expert_indices.repeat(1, orig_shape[1], 1) #(B, L, #num_experts_per)        
        
        x = x.view(orig_shape[0], orig_shape[1], 1, -1)  # (B, L, 1, C)
        x = x.repeat_interleave(self.num_experts_per, dim=2) # (B, L, num_experts_per, C)
        y = torch.empty_like(x, dtype=x.dtype) #(B, L, num_experts_per, C)
        for i, expert in enumerate(self.experts):
            y[expert_indices == i] = expert(x[expert_indices == i]).to(x.dtype)            
        
        y = (y.view(*expert_weights.shape, -1) * expert_weights.unsqueeze(-1)).sum(dim=2)
        x = x_orig + self.drop_path(y.view(*orig_shape))

        return x, scores



class PatchEmbed(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, sig_len=3000, conv_patch_sizes=[10,5,1], in_chans=1, embed_dim=[32,64,128], norm_layer=nn.LayerNorm):
        super().__init__()
        
        # num_patches = sig_len // conv_patch_size
        self.sig_len = sig_len
        self.conv_patch_sizes = conv_patch_sizes
        # self.num_patches = num_patches

        self.proj = nn.Conv1d(in_chans, embed_dim[0], kernel_size=self.conv_patch_sizes[0], stride=self.conv_patch_sizes[0])
        self.norm = norm_layer(embed_dim[0])
        self.act = nn.GELU()

        self.proj2 = nn.Conv1d(embed_dim[0], embed_dim[1], kernel_size=self.conv_patch_sizes[1], stride=self.conv_patch_sizes[1])
        self.norm2 = norm_layer(embed_dim[1])
        # self.act2 = nn.GELU()
        
        self.proj3 = nn.Conv1d(embed_dim[1], embed_dim[2], kernel_size=self.conv_patch_sizes[2], stride=self.conv_patch_sizes[2])
        self.norm3 = norm_layer(embed_dim[2])
        # self.act2 = nn.GELU()


    def forward(self, x):
        B, L, C = x.shape
        # # FIXME look at relaxing size constraints
        # assert L == self.sig_len, \
        #     f"Input Signal length ({L}) doesn't match model ({self.sig_len})."
        
        x = x.permute(0, 2, 1) # (B, C, L)
        x = self.proj(x)
        x = x.permute(0, 2, 1) # (B, L, C)

        x = self.act(self.norm(x))        

        x = x.permute(0, 2, 1) # (B, C, L)
        x = self.proj2(x)
        x = x.permute(0, 2, 1) # (B, L, C)

        x = self.act(self.norm2(x))        

        x = x.permute(0, 2, 1) # (B, C, L)
        x = self.proj3(x)
        x = x.permute(0, 2, 1) # (B, L, C)

        x = self.act(self.norm3(x))
        

        return x

class ConvNextBlock(nn.Module):
    """ ConvNeXt Block
    """
    def __init__(self, dim, out_dim, kernel_size=7, drop_path=0.3, stride=1, layer_scale_init_value=1e-6, norm_layer=nn.LayerNorm):
        super().__init__()
        self.pwconv0 = nn.Linear(dim, out_dim)
        self.dwconv = nn.Conv1d(out_dim, out_dim, kernel_size=kernel_size, stride=stride, groups=dim)  # depthwise conv
        self.norm = norm_layer(out_dim)
        self.pwconv1 = nn.Linear(out_dim, 4 * out_dim)  # pointwise/1x1 convs, implemented with linear layers
        self.act = nn.GELU()
        self.pwconv2 = nn.Linear(4 * out_dim, out_dim)
        self.gamma = nn.Parameter(layer_scale_init_value * torch.ones((out_dim)),
                                  requires_grad=True) if layer_scale_init_value > 0 else None
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.stride = stride

    def forward(self, x):
        
        x = x.permute(0, 2, 1)
        x = self.pwconv0(x)   
        x = x.permute(0, 2, 1)        
        x = self.dwconv(x)
        _x = x     
        x = x.permute(0, 2, 1)  # (B, L, C)        
        x = self.norm(x)
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        if self.gamma is not None:
            x = self.gamma.unsqueeze(0).unsqueeze(0) * x        
        x = x.permute(0, 2, 1)  # (B, C, L)
        x = _x + self.drop_path(x)
        # if self.stride == 1:
        #     x = input + self.drop_path(x)
            # input = input[:, :, ::self.stride]
        
        # x = input + self.drop_path(x)
        return x

class PatchUpsample(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, embed_dim=256, up_dim=128, upsample=2, norm_layer=nn.LayerNorm):
        super().__init__()

        self.upsample = upsample
        self.norm = norm_layer(embed_dim + up_dim)
        self.merge = nn.Linear(embed_dim + up_dim, up_dim)        
        
    def forward(self, x1, x2):
        """
        x1: (B, L, C)   From higher layer
        x2: (B, 2L, C)  From lower layer
        """

        x1 = x1.repeat_interleave(self.upsample, dim=1)  # (B, downsample*L, C)        
        
        x = torch.cat([x1,x2], dim=-1)
        x = self.norm(x)
        x = self.merge(x)        

        return x
    
class PatchMergeConv(nn.Module):
    """ Image to Patch Embedding
    """
    def __init__(self, embed_dim=128, exp_dim=256, downsample=2, norm_layer=nn.LayerNorm, drop_path = 0.1) :
        super().__init__()

        self.downsample = downsample
        self.exp = nn.Conv1d(embed_dim, exp_dim, kernel_size=1)
        kernel_size = downsample * 2
        padding = math.ceil(downsample / 2)
        self.dwconv = nn.Conv1d(exp_dim, exp_dim, kernel_size=kernel_size, groups=exp_dim, stride=downsample, padding=padding)        
        
        self.norm = norm_layer(exp_dim)
        self.pwconv1 = nn.Linear(exp_dim, exp_dim * 4)
        self.pwconv2 = nn.Linear(exp_dim * 4, exp_dim)
        self.act = nn.GELU()
        
        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        
    def forward(self, x):
        
        x = x.permute(0,2,1)
        x = self.exp(x)
        x = self.dwconv(x)        
        x = x.permute(0,2,1)        
        _x = x
        x = self.norm(x)        
        x = self.pwconv1(x)
        x = self.act(x)
        x = self.pwconv2(x)
        x = _x + self.drop_path(x)
        # x = self.act(x)

        return x


class RevIN(nn.Module):
    def __init__(self, num_features: int, eps=1e-8, affine=True):
        """
        :param num_features: the number of features or channels
        :param eps: a value added for numerical stability
        :param affine: if True, RevIN has learnable affine parameters
        """
        # super(RevIN, self).__init__()
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.affine = affine
        if self.affine:
            self._init_params()

    def forward(self, x, mask=None, mode:str='norm'):
        if mode == 'norm':
            self._get_statistics(x, mask)
            x = self._normalize(x)
        elif mode == 'denorm':
            x = self._denormalize(x)
        else: raise NotImplementedError
        return x

    def _init_params(self):
        # initialize RevIN params: (C,)
        self.affine_weight = nn.Parameter(torch.ones(self.num_features))
        self.affine_bias = nn.Parameter(torch.zeros(self.num_features))

    def _get_statistics(self, x, mask=None):                
        dim2reduce = tuple(range(1, x.ndim-1))
        # if mask is not None:
        #     # mask = torch.where(mask==0)[0]        
        #     # x_ = x[mask]
        #     x_ = x[(1 - mask).bool()].reshape(x.shape[0], -1, 1)
        # else:
        #     x_ = x        
        x_ = x # No mask

        self.mean = torch.mean(x_, dim=dim2reduce, keepdim=True).detach()
        self.stdev = torch.sqrt(torch.var(x_, dim=dim2reduce, keepdim=True, unbiased=False) + self.eps).detach()
        

    def _normalize(self, x):
        x = x - self.mean
        x = x / self.stdev
        if self.affine:
            x = x * self.affine_weight
            x = x + self.affine_bias
        return x

    def _denormalize(self, x):
        if self.affine:
            x = x - self.affine_bias
            x = x / (self.affine_weight + self.eps*self.eps)
        x = x * self.stdev
        x = x + self.mean
        return x

class RMSNorm(nn.Module):
    def __init__(self, normalized_shape, eps=1e-8):
        super().__init__()
        self.eps = eps
        self.gamma = nn.Parameter(torch.ones(normalized_shape))

    def forward(self, x):
        rms = torch.sqrt(torch.mean(x**2, dim=-1, keepdim=True))
        x_norm = x / (rms + self.eps)
        return self.gamma * x_norm
    
# Attention pooling from SleepFM
# https://github.com/zou-group/sleepfm-clinical/blob/main/sleepfm/models/models.py
class AttentionPooling(nn.Module):
    def __init__(self, input_dim, num_heads=4, norm_layer=nn.LayerNorm):
        super().__init__()
        self.transformer_layer = BlockFlashAttn(dim=input_dim, num_heads=num_heads, mlp_ratio=4.0, qkv_bias=True, norm_layer=norm_layer, drop_path=0.0)

    def forward(self, x):
        
        out, _ = self.transformer_layer(x)
        pooled_output = out.mean(dim=1)

        return pooled_output


def freq_pos_enc(dim, max_seq_len):
    # inverse frequencies [theta_1, theta_2, ..., theta_dim/2]
    inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2) / dim))  # (dim/2,)
    # positions up to max_seq_len
    pos = torch.arange(max_seq_len)  # (max_seq_len,)
    # frequency position encoding (outer product: pos x inv_freq)
    pos_enc = torch.einsum("p, f -> p f", pos, inv_freq)  # (max_seq_len, dim/2)
    # duplicate each element along inverse frequency dimension
    return repeat(pos_enc, "... f -> ... (f r)", r=2)  # (max_seq_len, dim)

def rotate_half(u):
    # last dimension of u is [u1, u2, u3, u4, ...]
    u = rearrange(u, '... (d r) -> ... d r', r=2)
    u1, u2 = u.unbind(dim=-1)
    u = torch.stack((-u2, u1), dim=-1)
    # last dimension of result is [-u2, u1, -u4, u3, ...]
    return rearrange(u, '... d r -> ... (d r)')


def rotate(u, pos_enc):
    num_tokens = u.shape[-2]
    # pos_enc = pos_enc[:num_tokens]
    return u * pos_enc.cos() + (rotate_half(u) * pos_enc.sin())