import torch
import torch.nn as nn
import torch.nn.functional as F

class MeshConv(nn.Module):
    """ Computes convolution between edges and 4 incident (1-ring) edge neighbors
    in the forward pass takes:
    x: edge features (Batch x Features x Edges)
    mesh: list of mesh data-structure (len(mesh) == Batch)
    and applies convolution
    """
    def __init__(self, in_channels, out_channels, k=5, bias=True):
        super(MeshConv, self).__init__()
        self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(1, k), bias=bias)
        self.k = k

    def forward(self, x, mesh):
        x = x.squeeze(-1)
        # pad gemm
        G = torch.cat([self.pad_gemm(i, x.shape[2], x.device) for i in mesh], 0)
        # build 'neighborhood image' and apply convolution
        G = self.create_GeMM(x, G)
        x = self.conv(G)
        return x

    def flatten_gemm_inds(self, Gi):
        (b, ne, nn) = Gi.shape
        ne += 1
        batch_n = torch.floor(torch.arange(b * ne, device=Gi.device).float() / ne).view(b, ne)
        add_fac = batch_n * ne
        add_fac = add_fac.view(b, ne, 1)
        add_fac = add_fac.repeat(1, 1, nn)
        # flatten Gi
        Gi = Gi.float() + add_fac[:, 1:, :]
        return Gi

    def create_GeMM(self, x, Gi):
        """ gathers the edge features (x) with from the 1-ring indices (Gi)
        applys symmetric functions to handle order invariance
        returns a 'fake image' which can use 2d convolution on
        output dimensions: Batch x Channels x Edges x 5
        """
        Gishape = Gi.shape
        # pad the first row of  every sample in batch with zeros
        padding = torch.zeros((x.shape[0], x.shape[1], 1), requires_grad=True, device=x.device)
        # add zero feature vector then shift all indices. border edges now reference zero vector
        x = torch.cat((padding, x), dim=2)
        Gi = Gi + 1 #shift

        # first flatten indices
        Gi_flat = self.flatten_gemm_inds(Gi)
        Gi_flat = Gi_flat.view(-1).long()
        #
        odim = x.shape
        x = x.permute(0, 2, 1).contiguous()
        x = x.view(odim[0] * odim[2], odim[1])

        # indices of gemm never reference padded section of x so padded section never used
        f = torch.index_select(x, dim=0, index=Gi_flat)
        f = f.view(Gishape[0], Gishape[1], Gishape[2], -1)
        f = f.permute(0, 3, 1, 2)

        # apply the symmetric functions for an equivariant convolution
        x_1 = f[:, :, :, 1] + f[:, :, :, 3]
        x_2 = f[:, :, :, 2] + f[:, :, :, 4]
        x_3 = torch.abs(f[:, :, :, 1] - f[:, :, :, 3])
        x_4 = torch.abs(f[:, :, :, 2] - f[:, :, :, 4])
        f = torch.stack([f[:, :, :, 0], x_1, x_2, x_3, x_4], dim=3)
        return f

    def pad_gemm(self, m, xsz, device):
        """ extracts one-ring neighbors (4x) -> m.gemm_edges
        which is of size #edges x 4
        add the edge_id itself to make #edges x 5
        then pad to desired size e.g., xsz x 5
        """
        padded_gemm = torch.tensor(m.gemm_edges, device=device).float()
        padded_gemm = padded_gemm.requires_grad_()
        padded_gemm = torch.cat((torch.arange(int(m.edges_count), device=device).float().unsqueeze(1), padded_gemm), dim=1)
        # pad using F
        padded_gemm = F.pad(padded_gemm, (0, 0, 0, xsz - m.edges_count), "constant", 0)
        padded_gemm = padded_gemm.unsqueeze(0)
        return padded_gemm
[((14, 20, 14, 112), 'torch.nn.Conv2d', 'nn.Conv2d', (), '', True, 'import torch.nn as nn\n'), ((45, 18, 45, 95), 'torch.zeros', 'torch.zeros', (), '', False, 'import torch\n'), ((47, 12, 47, 42), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((59, 12, 59, 55), 'torch.index_select', 'torch.index_select', (), '', False, 'import torch\n'), ((66, 14, 66, 54), 'torch.abs', 'torch.abs', ({(66, 24, 66, 53): 'f[:, :, :, (1)] - f[:, :, :, (3)]'}, {}), '(f[:, :, :, (1)] - f[:, :, :, (3)])', False, 'import torch\n'), ((67, 14, 67, 54), 'torch.abs', 'torch.abs', ({(67, 24, 67, 53): 'f[:, :, :, (2)] - f[:, :, :, (4)]'}, {}), '(f[:, :, :, (2)] - f[:, :, :, (4)])', False, 'import torch\n'), ((68, 12, 68, 67), 'torch.stack', 'torch.stack', (), '', False, 'import torch\n'), ((81, 22, 81, 87), 'torch.nn.functional.pad', 'F.pad', ({(81, 28, 81, 39): 'padded_gemm', (81, 41, 81, 71): '(0, 0, 0, xsz - m.edges_count)', (81, 73, 81, 83): '"""constant"""', (81, 85, 81, 86): '0'}, {}), "(padded_gemm, (0, 0, 0, xsz - m.edges_count), 'constant', 0)", True, 'import torch.nn.functional as F\n'), ((77, 22, 77, 63), 'torch.tensor', 'torch.tensor', (), '', False, 'import torch\n'), ((29, 30, 29, 68), 'torch.arange', 'torch.arange', (), '', False, 'import torch\n')]