# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/nets/10_symswin3d.ipynb.

# %% auto 0
__all__ = ['symmetry_attention_rearrange_forward', 'symmetry_attention_rearrange_backward']

# %% ../../nbs/nets/10_symswin3d.ipynb 6
import torch
from einops import rearrange

# %% ../../nbs/nets/10_symswin3d.ipynb 8
def symmetry_attention_rearrange_forward(hidden_states: torch.Tensor):
    # hidden_states: (b, num_patches_z, num_patches_y, num_patches_x, dim)

    # Flip the second half of the patches along the x axis
    last_dim_length = hidden_states.shape[-2]
    hidden_states = torch.cat(
        [
            hidden_states[..., : last_dim_length // 2, :],
            hidden_states[..., last_dim_length // 2 :, :].flip(-2),
        ],
        dim=-2,
    )
    # (b, num_patches_z, num_patches_y, num_patches_x_first_half + flipped_num_patches_second_half, dim)

    # Rearrange the patches to alternate between the two halves along the x axis
    hidden_states = rearrange(
        hidden_states,
        "b num_patches_z num_patches_y (two half_num_patches_x) dim -> "
        "b num_patches_z num_patches_y (half_num_patches_x two) dim",
        two=2,
    ).contiguous()
    # (b, num_patches_z, num_patches_y, rearranged_num_patches_x, dim)

    return hidden_states


def symmetry_attention_rearrange_backward(hidden_states: torch.Tensor):
    # hidden_states: (b, num_patches_z, num_patches_y, rearranged_num_patches_x, dim)

    # Return the patches to their previous order along the x-axis
    hidden_states = rearrange(
        hidden_states,
        "b num_patches_z num_patches_y (half_num_patches_x two) dim -> "
        "b num_patches_z num_patches_y (two half_num_patches_x) dim",
        two=2,
    ).contiguous()
    # (b, num_patches_z, num_patches_y, num_patches_x_first_half + flipped_num_patches_second_half, dim)

    # Flip the second half of the patches along the x axis to return to the original state
    last_dim_length = hidden_states.shape[-2]
    hidden_states = torch.cat(
        [
            hidden_states[..., : last_dim_length // 2, :],
            hidden_states[..., last_dim_length // 2 :, :].flip(-2),
        ],
        dim=-2,
    )
    # (b, num_patches_z, num_patches_y, num_patches_x, dim)

    return hidden_states
