import torch
from torch import nn
from torch.nn.parameter import Parameter

from einops import rearrange, reduce, repeat

class dca_offsets_layer(nn.Module):
    """Constructs a Offset Generation module.
    """
    def __init__(self, channel, n_offsets):
        super(dca_offsets_layer, self).__init__()

        self.channel = channel
        self.n_offsets = n_offsets

    def covariance_features(self, x):
        """
        Takes in a feature map and returns the unnormalized covariance matrix
        """
        m_batchsize, C, height, width = x.size()
        x = x - x.mean(dim=1, keepdim=True) / (x.std(dim=1, keepdim=True) + 1e-5)
        proj_query = x.view(m_batchsize, C, -1)
        proj_key = x.view(m_batchsize, C, -1).permute(0, 2, 1)
        energy = torch.bmm(proj_query, proj_key)
        return energy

    def forward(self, x):
        m_batchsize, C, height, width = x.size()
        cov_matrix = self.covariance_features(x).reshape(m_batchsize, C, 1, C)

        _, locations = torch.topk(cov_matrix, self.n_offsets, dim=1)
        delta = torch.stack(self.n_offsets*[torch.arange(0, self.channel)], dim=0)
        delta = torch.stack(m_batchsize * [delta], dim=0)
        offsets = locations.squeeze() - delta.cuda()
        return offsets 

[((24, 17, 24, 48), 'torch.bmm', 'torch.bmm', ({(24, 27, 24, 37): 'proj_query', (24, 39, 24, 47): 'proj_key'}, {}), '(proj_query, proj_key)', False, 'import torch\n'), ((31, 23, 31, 68), 'torch.topk', 'torch.topk', (), '', False, 'import torch\n'), ((33, 16, 33, 57), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((32, 44, 32, 73), 'torch.arange', 'torch.arange', ({(32, 57, 32, 58): '0', (32, 60, 32, 72): 'self.channel'}, {}), '(0, self.channel)', False, 'import torch\n')]