import torch


def permutation_kernel(dim):
    """
    Outputs the kernel for a permutation-based invertible downsampling layer.
    The same kernel can be used for the invertible upsampling, using the 
    transposed convolution. Works for any data format (1D, 2D, 3D, 543D,...)
    """
    
    channel_factor = 2**dim
    
    kernel_atom_shape = ()
    for i in range(dim):
        kernel_atom_shape+= (2,)
    
    # We first create zero-kernel-atoms of the appropriate size and then
    # fill a 1 at the correct place. The index is determined with the help
    # of a binary representation, e.g. '10' results in a 1 in position [1][0].
    kernel_atoms = []
    for i in range(channel_factor):
        kernel_atom = torch.zeros(kernel_atom_shape)
        indices_bitstring = ('{0:0'+str(dim)+'b}').format(i)
        indices = ()
        for j in range(len(indices_bitstring)):
            indices+= (int(indices_bitstring[j]),)
        kernel_atom[indices] = 1
        kernel_atoms.append(kernel_atom)
        
    kernel = torch.stack(kernel_atoms, dim=0).view(channel_factor,1,*([2]*dim))
    return kernel