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, 'torch.nn.Conv2d', 'nn.Conv2d', ([], {'in_channels': 'in_channels', 'out_channels': 'out_channels', 'kernel_size': '(1, k)', 'bias': 'bias'}), True, 'import torch.nn as nn\n'), (45, 'torch.zeros', 'torch.zeros', (['(x.shape[0], x.shape[1], 1)'], {'requires_grad': '(True)', 'device': 'x.device'}), False, 'import torch\n'), (47, 'torch.cat', 'torch.cat', (['(padding, x)'], {'dim': '(2)'}), False, 'import torch\n'), (59, 'torch.index_select', 'torch.index_select', (['x'], {'dim': '(0)', 'index': 'Gi_flat'}), False, 'import torch\n'), (66, 'torch.abs', 'torch.abs', (['(f[:, :, :, (1)] - f[:, :, :, (3)])'], {}), False, 'import torch\n'), (67, 'torch.abs', 'torch.abs', (['(f[:, :, :, (2)] - f[:, :, :, (4)])'], {}), False, 'import torch\n'), (68, 'torch.stack', 'torch.stack', (['[f[:, :, :, (0)], x_1, x_2, x_3, x_4]'], {'dim': '(3)'}), False, 'import torch\n'), (81, 'torch.nn.functional.pad', 'F.pad', (['padded_gemm', '(0, 0, 0, xsz - m.edges_count)', '"""constant"""', '(0)'], {}), True, 'import torch.nn.functional as F\n'), (77, 'torch.tensor', 'torch.tensor', (['m.gemm_edges'], {'device': 'device'}), False, 'import torch\n'), (29, 'torch.arange', 'torch.arange', (['(b * ne)'], {'device': 'Gi.device'}), False, 'import torch\n')]