import math
from collections import Counter

import numpy as np
from scipy import sparse
import torch
from torch.nn import functional as F
from torch.nn import Conv1d
from modules import remap
from sparselinear import SparseLinear


def reformat(x):
    """Reformat the input from a 4D tensor to a 3D tensor
    Args:
        x (:obj:`torch.tensor`): a 4D tensor
    Returns:
        :obj:`torch.tensor`: a 3D tensor
    """
    x = x.permute(0, 2, 3, 1)
    N, D1, D2, Feat = x.size()
    x = x.view(N, D1 * D2, Feat)
    return x


def equiangular_dimension_unpack(nodes, ratio):
    """Calculate the two underlying dimensions
    from the total number of nodes
    Args:
        nodes (int): combined dimensions
        ratio (float): ratio between the two dimensions
    Returns:
        int, int: separated dimensions
    """
    dim1 = int((nodes / ratio) ** 0.5)
    dim2 = int((nodes * ratio) ** 0.5)
    if dim1 * dim2 != nodes: # Try to correct dim1 or dim2 if ratio is wrong
        if nodes % dim1 == 0:
            dim2 = nodes // dim1
        if nodes % dim2 == 0:
            dim1 = nodes // dim2
    assert dim1 * dim2 == nodes, f'Unable to unpack nodes: {nodes}, ratio: {ratio}'
    return dim1, dim2


def equiangular_calculator(tensor, ratio):
    N, M, F = tensor.size()
    dim1, dim2 = equiangular_dimension_unpack(M, ratio)
    # bw_dim1, bw_dim2 = dim1 / 2, dim2 / 2
    tensor = tensor.view(N, dim1, dim2, F)
    return tensor


# 2D CNN layers
class Conv2dEquiangular(torch.nn.Module):
    """ Equiangular 2D Convolutional layer, periodic in the longitude (width) dimension.

    Parameters
    ----------
    in_channels : int
        Number of channels in the input image.
    out_channels : int
        Number of channels in the output image.
    kernel_size : int
        Width of the square convolutional kernel.
        The actual size of the kernel is kernel_size**2
    ratio : int
        ratio = H // W. Aspect ratio to reorganize the equiangular map from 1D to 2D
    periodic : bool
        whether to use periodic padding. (default: True)
    """

    def __init__(self, in_channels, out_channels, kernel_size, ratio, periodic, **kwargs):
        super().__init__()
        self.ratio = ratio
        self.periodic = periodic

        self.kernel_size = kernel_size
        self.pad_width = int((self.kernel_size - 1) / 2)

        self.conv = torch.nn.Conv2d(in_channels, out_channels, self.kernel_size)

        torch.nn.init.xavier_uniform_(self.conv.weight)
        torch.nn.init.zeros_(self.conv.bias)

    def periodicPad(self, x, width, periodic=True):
        """ Periodic padding function.

        Parameters
        ----------
        x : torch.Tensor
            The tensor format should be N * C * H * W
        width : int
            The padding width.
        periodic: bool
            Whether to use periodic padding. If False, the function is same as ZeroPad2D. (default: True)
        """
        if periodic:
            x = torch.cat((x[:, :, :, -width:], x, x[:, :, :, :width]), dim=3)
            padded = F.pad(x, (0, 0, width, width), 'constant', 0)
        else:
            padded = F.pad(x, (width, width, width, width), 'constant', 0)
        return padded

    def forward(self, x):
        x = equiangular_calculator(x, self.ratio)
        x = x.permute(0, 3, 1, 2)

        x = self.periodicPad(x, self.pad_width, self.periodic)
        x = self.conv(x)

        x = reformat(x)
        return x



# Graph CNN layers
"""
PyTorch implementation of a convolutional neural network on graphs based on
Chebyshev polynomials of the graph Laplacian.
See https://arxiv.org/abs/1606.09375 for details.
Copyright 2018 Michaël Defferrard.
Released under the terms of the MIT license.
"""


def prepare_laplacian(laplacian):
    """Prepare a graph Laplacian to be fed to a graph convolutional layer
    """

    def estimate_lmax(laplacian, tol=5e-3):
        r"""Estimate the largest eigenvalue of an operator."""
        lmax = sparse.linalg.eigsh(laplacian, k=1, tol=tol,
                                   ncv=min(laplacian.shape[0], 10),
                                   return_eigenvectors=False)
        lmax = lmax[0]
        lmax *= 1 + 2 * tol  # Be robust to errors.
        return lmax

    def scale_operator(L, lmax, scale=1):
        r"""Scale the eigenvalues from [0, lmax] to [-scale, scale]."""
        I = sparse.identity(L.shape[0], format=L.format, dtype=L.dtype)
        L *= 2 * scale / lmax
        L -= I
        return L

    lmax = estimate_lmax(laplacian)
    laplacian = scale_operator(laplacian, lmax)

    laplacian = sparse.coo_matrix(laplacian)

    # PyTorch wants a LongTensor (int64) as indices (it'll otherwise convert).
    indices = np.empty((2, laplacian.nnz), dtype=np.int64)
    np.stack((laplacian.row, laplacian.col), axis=0, out=indices)
    indices = torch.from_numpy(indices)

    laplacian = torch.sparse_coo_tensor(indices, laplacian.data, laplacian.shape)
    laplacian = laplacian.coalesce()  # More efficient subsequent operations.
    return laplacian


def cheb_conv(laplacian, inputs, weight):
    """Chebyshev convolution.
    Parameters
    ----------
    laplacian : torch.sparse.Tensor
        The laplacian corresponding to the current sampling of the sphere
    inputs : torch.Tensor
        The current input data being forwarded
    weight : torch.Tensor
        The weights of the current layer
        
    Returns
    -------
    x : torch.Tensor
        Inputs after applying Chebyshev convolution.
    """

    B, V, Fin1 = inputs.shape
    # print('B: {}, V. {}, Fin: {}'.format(B,V,Fin1))
    Fin, K, Fout = weight.shape
    # print('Fin: {}, K: {}, Fout: {}'.format(Fin, K, Fout))
    assert Fin1 == Fin
    # B = batch size
    # V = nb vertices
    # Fin = nb input features
    # Fout = nb output features
    # K = order of Chebyshev polynomials (kenel size)

    # transform to Chebyshev basis    
    x0 = inputs.permute(1, 2, 0).contiguous()  # V x Fin x B
    x0 = x0.view([V, Fin * B])  # V x Fin*B
    x = x0.unsqueeze(0)  # 1 x V x Fin*B

    if K > 1:
        x1 = torch.sparse.mm(laplacian, x0)  # V x Fin*B
        x = torch.cat((x, x1.unsqueeze(0)), 0)  # 2 x V x Fin*B
    for _ in range(2, K):
        x2 = 2 * torch.sparse.mm(laplacian, x1) - x0
        x = torch.cat((x, x2.unsqueeze(0)), 0)  # M x Fin*B
        x0, x1 = x1, x2

    x = x.view([K, V, Fin, B])  # K x V x Fin x B
    x = x.permute(3, 1, 2, 0).contiguous()  # B x V x Fin x K
    x = x.view([B * V, Fin * K])  # B*V x Fin*K

    # Linearly compose Fin features to get Fout features
    weight = weight.view(Fin * K, Fout)
    x = x.matmul(weight)  # B*V x Fout
    x = x.view([B, V, Fout])  # B x V x Fout

    return x


class ConvCheb(torch.nn.Module):
    """Graph convolutional layer.

    Parameters
    ----------
    in_channels : int
        Number of channels in the input graph.
    out_channels : int
        Number of channels in the output graph.
    kernel_size : int
        Number of trainable parameters per filter, which is also the size of
        the convolutional kernel.
        The order of the Chebyshev polynomials is kernel_size - 1.

        * A kernel_size of 1 won't take the neighborhood into account.
        * A kernel_size of 2 will look up to the 1-neighborhood (1 hop away).
        * A kernel_size of 3 will look up to the 2-neighborhood (2 hops away).

        A kernel_size of 0 is equivalent to not having a graph (or an empty
        adjacency matrix). All the vertices are treated independently and form
        a set. Every element of that set is given to a fully connected layer
        with a weight matrix of size (out_channels x in_channels).
    bias : bool
        Whether to add a bias term.
    conv : callable
        Function which will perform the actual convolution.
    """

    def __init__(self, in_channels, out_channels, kernel_size, laplacian, bias=True,
                 conv=cheb_conv, **kwargs):
        super().__init__()

        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.register_buffer(f'laplacian', laplacian)
        self._conv = conv

        # shape = (kernel_size, out_channels, in_channels)
        shape = (in_channels, kernel_size, out_channels)
        self.weight = torch.nn.Parameter(torch.Tensor(*shape))

        if bias:
            self.bias = torch.nn.Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)

        self.reset_parameters()

    def reset_parameters(self, activation='relu', fan='in',
                         distribution='normal'):
        r"""Reset weight and bias.

        * Kaiming / He is given by `activation='relu'`, `fan='in'` or
        `fan='out'`, and `distribution='normal'`.
        * Xavier / Glorot is given by `activation='linear'`, `fan='avg'`,
          and `distribution='uniform'`.
        * LeCun is given by `activation='linear'` and `fan='in'`.

        Motivation based on inits from PyTorch, TensorFlow, Keras.

        Parameters
        ----------
        activation : {'relu', 'linear', 'sigmoid', 'tanh'}
            Select the activation function your are using.
        fan : {'in', 'out', 'avg'}
            Select `'in'` to preserve variance in the forward pass.
            Select `'out'` to preserve variance in the backward pass.
            Select `'avg'` for a balance.
        distribution : {'normal', 'uniform'}
            Whether to draw weights from a normal or random distribution.

        References
        ----------
        Delving Deep into Rectifiers: Surpassing Human-Level Performance on
        ImageNet Classification, Kaiming He, Xiangyu Zhang, Shaoqing Ren, Jian
        Sun, https://arxiv.org/abs/1502.01852

        Understanding the difficulty of training deep feedforward neural
        networks, Xavier Glorot, Yoshua Bengio,
        http://proceedings.mlr.press/v9/glorot10a.html
        """

        if fan == 'in':
            fan = self.in_channels * self.kernel_size
        elif fan == 'out':
            fan = self.out_channels * self.kernel_size
        elif fan == 'avg':
            fan = (self.in_channels + self.out_channels) / 2 * self.kernel_size
        else:
            raise ValueError('unknown fan')

        if activation == 'relu':
            scale = 2  # relu kills half the activations, from He et al.
        elif activation in ['linear', 'sigmoid', 'tanh']:
            # sigmoid and tanh are linear around 0
            scale = 1  # from Glorot et al.
        else:
            raise ValueError('unknown activation')

        if distribution == 'normal':
            std = math.sqrt(scale / fan)
            self.weight.data.normal_(0, std)
        elif distribution == 'uniform':
            limit = math.sqrt(3 * scale / fan)
            self.weight.data.uniform_(-limit, limit)
        else:
            raise ValueError('unknown distribution')

        if self.bias is not None:
            self.bias.data.fill_(0)

    def set_parameters(self, weight, bias=None):
        r"""Set weight and bias.

        Parameters
        ----------
        weight : array of shape in_channels x kernel_size x out_channels
            The coefficients of the Chebyshev polynomials.
        bias : vector of length out_channels
            The bias.
        """
        self.weight = torch.nn.Parameter(torch.as_tensor(weight))
        if bias is not None:
            self.bias = torch.nn.Parameter(torch.as_tensor(bias))

    def extra_repr(self):
        s = '{in_channels} -> {out_channels}, kernel_size={kernel_size}'
        s += ', bias=' + str(self.bias is not None)
        return s.format(**self.__dict__)

    def forward(self, inputs):
        r"""Forward graph convolution.

        Parameters
        ----------
        laplacian : sparse matrix of shape n_vertices x n_vertices
            Encode the graph structure.
        inputs : tensor of shape n_signals x n_vertices x n_features
            Data, i.e., features on the vertices.
        """

        outputs = self._conv(self.laplacian, inputs, self.weight)
        if self.bias is not None:
            outputs += self.bias
        return outputs


# Conv layers
class Conv1dAuto(Conv1d):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.padding = (self.kernel_size[0] // 2)  # dynamic add padding based on the kernel_size



class PoolMaxEquiangular(torch.nn.MaxPool1d):
    """EquiAngular max pooling module
    
    Parameters
    ----------
    ratio : float
        Ratio between latitude and longitude dimensions of the data
    kernel_size : int
        Pooling kernel width
    return_indices : bool (default : True)
        Whether to return the indices corresponding to the locations of the maximum value retained at pooling
    """

    def __init__(self, ratio, kernel_size, return_indices=True, *args, **kwargs):
        self.ratio = ratio
        kernel_size = int(kernel_size ** 0.5)
        super().__init__(kernel_size=kernel_size, return_indices=return_indices)

    def forward(self, inputs):
        """calls Maxpool1d and if desired, keeps indices of the pixels pooled to unpool them
        Parameters
        ----------
        x : torch.tensor of shape batch x pixels x features
            Input data
            
        Returns
        -------
        x : torch.tensor of shape batch x unpooled pixels x features
            Layer output
        indices : list(int)
            Indices of the pixels pooled
        """
        x = equiangular_calculator(inputs, self.ratio)
        x = x.permute(0, 3, 1, 2)

        if self.return_indices:
            x, indices = F.max_pool2d(x, self.kernel_size, return_indices=self.return_indices)
        else:
            x = F.max_pool2d(x, self.kernel_size)
        x = reformat(x)

        if self.return_indices:
            output = x, indices
        else:
            output = x

        return output


class UnpoolMaxEquiangular(torch.nn.MaxUnpool1d):
    """Equiangular max unpooling module
    
    Parameters
    ----------
    ratio : float
        Ratio between latitude and longitude dimensions of the data
    kernel_size : int
        Pooling kernel width
    """

    def __init__(self, ratio, kernel_size, *args, **kwargs):
        self.ratio = ratio
        kernel_size = int(kernel_size ** 0.5)
        super().__init__(kernel_size=(kernel_size, kernel_size))

    def forward(self, inputs, indices):
        """calls MaxUnpool1d using the indices returned previously by PoolMaxEquiangular
        Parameters
        ----------
        inputs : torch.tensor of shape batch x pixels x features
            Input data
        indices : int
            Indices of pixels equiangular maxpooled previously
            
        Returns
        -------
        x : torch.tensor of shape batch x unpooled pixels x features
            Layer output
        """
        x = equiangular_calculator(inputs, self.ratio)
        x = x.permute(0, 3, 1, 2)
        x = F.max_unpool2d(x, indices, self.kernel_size)
        x = reformat(x)
        return x


class PoolAvgEquiangular(torch.nn.AvgPool1d):
    """EquiAngular average pooling
    
    Parameters
    ----------
    ratio : float
        Parameter for equiangular sampling -> width/height
    kernel_size : int
        Pooling kernel width
    """

    def __init__(self, ratio, kernel_size, *args, **kwargs):
        self.ratio = ratio
        kernel_size = int(kernel_size ** 0.5)
        super().__init__(kernel_size=(kernel_size, kernel_size))

    def forward(self, inputs):
        """calls Avgpool1d
        Parameters
        ----------
        inputs : torch.tensor of shape batch x pixels x features
            Input data
        
        Returns
        -------
        x : torch.tensor of shape batch x pooled pixels x features
            Layer output
        """
        x = equiangular_calculator(inputs, self.ratio)
        x = x.permute(0, 3, 1, 2)
        x = F.avg_pool2d(x, self.kernel_size)
        x = reformat(x)

        return x, None


class UnpoolAvgEquiangular(torch.nn.Module):
    """EquiAngular average unpooling
    
    Parameters
    ----------
    ratio : float
        Parameter for equiangular sampling -> width/height
    """

    def __init__(self, ratio, kernel_size, *args, **kwargs):
        self.ratio = ratio
        self.kernel_size = int(kernel_size ** 0.5)
        super().__init__()

    def forward(self, inputs, *args):
        """calls pytorch's interpolate function to create the values while unpooling based on the nearby values
        Parameters
        ----------
        inputs : torch.tensor of shape batch x pixels x features
            Input data
        
        Returns
        -------
        x : torch.tensor of shape batch x unpooled pixels x features
            Layer output
        """

        x = equiangular_calculator(inputs, self.ratio)
        x = x.permute(0, 3, 1, 2)
        x = F.interpolate(x, scale_factor=(self.kernel_size, self.kernel_size), mode="nearest")
        x = reformat(x)
        return x


class PoolMaxHealpix(torch.nn.MaxPool1d):
    """Healpix Maxpooling module
     
    Parameters
    ----------
    kernel_size : int
        Pooling kernel width
    return_indices : bool (default : True)
        Whether to return the indices corresponding to the locations of the maximum value retained at pooling
    """

    def __init__(self, kernel_size, return_indices=True, *args, **kwargs):
        super().__init__(kernel_size=kernel_size, return_indices=return_indices)

    def forward(self, x):
        """calls Maxpool1d and if desired, keeps indices of the pixels pooled to unpool them
        Parameters
        ----------
        x : torch.tensor of shape batch x pixels x features
            Input data  
        indices : list
            Indices where the max value was located in unpooled image
            
        Returns
        -------
        x : torch.tensor of shape batch x unpooled pixels x features
            Layer output
        indices : list(int)
            Indices of the pixels pooled
        """
        x = x.permute(0, 2, 1)
        if self.return_indices:
            x, indices = F.max_pool1d(x, self.kernel_size, return_indices=self.return_indices)
        else:
            x = F.max_pool1d(x)
        x = x.permute(0, 2, 1)

        if self.return_indices:
            output = x, indices
        else:
            output = x
        return output


class PoolAvgHealpix(torch.nn.Module):
    """Healpix average pooling module
    
    Parameters
    ----------
    kernel_size : int
        Pooling kernel width
    """

    def __init__(self, kernel_size, *args, **kwargs):
        """kernel_size should be 4, 16, 64, etc."""
        super().__init__()
        self.kernel_size = kernel_size

    def extra_repr(self):
        return 'kernel_size={kernel_size}'.format(**self.__dict__)

    def forward(self, x):
        """x has shape (batch, pixels, channels) and is in nested ordering"""
        x = x.permute(0, 2, 1)
        x = torch.nn.functional.avg_pool1d(x, self.kernel_size)
        return x.permute(0, 2, 1), None


class UnpoolAvgHealpix(torch.nn.Module):
    """Healpix Average Unpooling module
    
    Parameters
    ----------
    kernel_size : int
        Pooling kernel width
    """

    def __init__(self, kernel_size, *args, **kwargs):
        """kernel_size should be 4, 16, 64, etc."""
        super().__init__()
        self.kernel_size = kernel_size

    def extra_repr(self):
        return 'kernel_size={kernel_size}'.format(**self.__dict__)

    def forward(self, x, *args):
        """x has shape (batch, pixels, channels) and is in nested ordering"""
        # return x.repeat_interleave(self.kernel_size, dim=1)
        x = x.permute(0, 2, 1)
        x = torch.nn.functional.interpolate(x, scale_factor=self.kernel_size, mode='nearest')
        return x.permute(0, 2, 1)


class UnpoolMaxHealpix(torch.nn.MaxUnpool1d):
    """HEALpix max unpooling module
    
    Parameters
    ----------
    kernel_size : int
        Pooling kernel width
    """

    def __init__(self, kernel_size, *args, **kwargs):
        super().__init__(kernel_size=kernel_size)

    def forward(self, x, indices, **kwargs):
        """calls pytorch's unpool1d function to create the values while unpooling based on the nearby values
        Parameters
        ----------
        **kwargs
        inputs : torch.tensor of shape batch x pixels x features
            Input data
        indices : list
            Indices where the max value was located in unpooled image
        
        Returns
        -------
        x : torch.tensor of shape batch x unpooled pixels x features
            Layer output
        """

        x = x.permute(0, 2, 1)
        x = F.max_unpool1d(x, indices, self.kernel_size)
        x = x.permute(0, 2, 1)
        return x


def _build_interpolation_matrix(src_graph, dst_graph):
    """Return the sparse matrix that interpolates between two spherical samplings."""

    ds = remap.compute_interpolation_weights(src_graph, dst_graph, method='conservative', normalization='fracarea') # destarea’

    # Sanity checks.
    np.testing.assert_allclose(ds.src_grid_center_lat, src_graph.signals['lat'])
    np.testing.assert_allclose(ds.src_grid_center_lon, src_graph.signals['lon'])
    np.testing.assert_allclose(ds.dst_grid_center_lat, dst_graph.signals['lat'])
    np.testing.assert_allclose(ds.dst_grid_center_lon, dst_graph.signals['lon'])
    np.testing.assert_allclose(ds.src_grid_frac, 1)
    np.testing.assert_allclose(ds.dst_grid_frac, 1)
    np.testing.assert_allclose(ds.src_grid_imask, 1)
    np.testing.assert_allclose(ds.dst_grid_imask, 1)

    col = ds.src_address
    row = ds.dst_address
    dat = ds.remap_matrix.squeeze()
    # CDO indexing starts at 1
    row = np.array(row) - 1
    col = np.array(col) - 1
    weights = sparse.csr_matrix((dat, (row, col)))
    assert weights.shape == (dst_graph.n_vertices, src_graph.n_vertices)

    # Destination pixels are normalized to 1 (row-sum = 1).
    # Weights represent the fractions of area attributed to source pixels.
    np.testing.assert_allclose(weights.sum(axis=1), 1)
    # Interpolation is conservative: it preserves area.
    np.testing.assert_allclose(weights.T @ ds.dst_grid_area, ds.src_grid_area)

    # Unnormalize.
    weights = weights.multiply(ds.dst_grid_area.values[:, np.newaxis])

    # Another way to assert that the interpolation is conservative.
    np.testing.assert_allclose(np.asarray(weights.sum(1)).squeeze(), ds.dst_grid_area)
    np.testing.assert_allclose(np.asarray(weights.sum(0)).squeeze(), ds.src_grid_area)

    return weights


def build_pooling_matrices(src_graph, dst_graph):
    weights = _build_interpolation_matrix(src_graph, dst_graph)
    pool = weights.multiply(1/weights.sum(1))
    unpool = weights.multiply(1/weights.sum(0)).T
    return pool, unpool


def convert_to_torch_sparse(mat: "sparse.coo.coo_matrix"):
    indices = np.empty((2, mat.nnz), dtype=np.int64)
    np.stack((mat.row, mat.col), axis=0, out=indices)
    indices = torch.from_numpy(indices)
    mat = torch.sparse_coo_tensor(indices, mat.data, mat.shape, dtype=torch.float32)
    mat = mat.coalesce()
    return mat


class RemapBlock(torch.nn.Module):
    def __init__(self, remap_matrix: "sparse.coo.coo_matrix"):
        super().__init__()
        remap_matrix = self.process_remap_matrix(remap_matrix)
        self.register_buffer('remap_matrix', remap_matrix)
    
    def forward(self, x, *args, **kwargs):
        n_batch, n_nodes, n_val = x.shape
        matrix = self.remap_matrix
        new_nodes, _ = matrix.shape
        x = x.permute(1, 2, 0).reshape(n_nodes, n_batch * n_val)
        x = torch.sparse.mm(matrix, x)
        x = x.reshape(new_nodes, n_val, n_batch).permute(2, 0, 1)
        return x
    
    def process_remap_matrix(self, mat):
        return convert_to_torch_sparse(mat)


class GeneralAvgPool(RemapBlock):
    def forward(self, x, *args, **kwargs):
        x = super().forward(x, *args, **kwargs)
        # Some pooling methods (e.g. Avg) do not give source indices, please return None for compatibility reason
        return x, None


class GeneralAvgUnpool(RemapBlock):
    def forward(self, x, *args, **kwargs):
        x = super().forward(x, *args, **kwargs)
        return x


class GeneralMaxAreaPool(RemapBlock):
    def forward(self, x, *args, **kwargs):
        x = super().forward(x, *args, **kwargs)
        # Some pooling methods (e.g. Avg) do not give source indices, please return None for compatibility reason
        return x, None
        
    def process_remap_matrix(self, mat):
        max_ind_col = np.argmax(mat, axis=1).T
        row = np.arange(max_ind_col.shape[1]).reshape(1, -1)
        indices = np.concatenate([row, max_ind_col], axis=0)
        indices = torch.from_numpy(indices.astype(np.int64))
        mat = torch.sparse_coo_tensor(indices, torch.ones(indices.shape[1]), mat.shape, dtype=torch.float32)
        mat = mat.coalesce()
        return mat


class GeneralMaxAreaUnpool(RemapBlock):
    def process_remap_matrix(self, mat):
        max_ind_row = np.argmax(mat, axis=0)
        col = np.arange(max_ind_row.shape[1]).reshape(1, -1)
        indices = np.concatenate([max_ind_row, col], axis=0)
        indices = torch.from_numpy(indices.astype(np.int64))
        mat = torch.sparse_coo_tensor(indices, torch.ones(indices.shape[1]), mat.shape, dtype=torch.float32)
        mat = mat.coalesce()
        return mat


class GeneralLearnableUnpool(SparseLinear):
    def __init__(self, remap_matrix: "sparse.coo.coo_matrix"):
        out_feature, in_feature = remap_matrix.shape
        bias = False
        indices = np.empty((2, remap_matrix.nnz), dtype=np.int64)
        np.stack((remap_matrix.row, remap_matrix.col), axis=0, out=indices)
        indices = torch.from_numpy(indices)
        dynamic = False
        super().__init__(in_feature, out_feature, bias, connectivity=indices, dynamic=dynamic)
    
    def forward(self, x, *args, **kwargs):
        x = x.permute(0, 2, 1)
        x = super().forward(x)
        x = x.permute(0, 2, 1)
        return x


class GeneralLearnablePool(GeneralLearnableUnpool):    
    def forward(self, inputs):
        output = super().forward(inputs)
        return output, None


class GeneralMaxValPool(RemapBlock):
    def forward(self, x, *args, **kwargs):
        n_batch, n_nodes, n_val = x.shape
        matrix = self.remap_matrix
        if not matrix.is_coalesced():
            matrix = matrix.coalesce()
        new_nodes, old_nodes = matrix.shape
        assert n_nodes == old_nodes, 'remap_matrix.shape[1] != input.shape[1]'
        x = x.permute(1, 2, 0).reshape(n_nodes, n_batch * n_val)

        indices = matrix.indices()
        row, col = indices
        weights = matrix.values()

        cnt = Counter([i.item() for i in row])
        kernel_sizes = [cnt[i] for i in sorted(cnt)]

        col = col.repeat(n_batch * n_val, 1).T

        val = torch.gather(x, dim=0, index=col).detach()
        val.requires_grad_(False)
        weighted_val = weights.view(-1, 1) * val

        start_row = 0
        max_val_index = []
        for k in kernel_sizes:
            curr = weighted_val[start_row:start_row+k]
            max_val_index.append(torch.argmax(curr, dim=0) + start_row)
            start_row += k
        max_val_index = torch.stack(max_val_index)
        nnz_row = torch.gather(col, dim=0, index=max_val_index)

        x_pooled = torch.gather(x, dim=0, index=nnz_row)

        _, nnz_col = np.indices(x_pooled.shape)
        nnz_col = torch.LongTensor(nnz_col).to(nnz_row.device)
        nnz_ind = torch.stack([nnz_row, nnz_col], dim=2)
        nnz_ind = nnz_ind.permute(1, 0, 2).reshape(-1, 2).T
        nnz_ind.requires_grad_(False)

        x_pooled = x_pooled.reshape(new_nodes, n_val, n_batch).permute(2, 0, 1)

        return x_pooled, nnz_ind


class GeneralMaxValUnpool(RemapBlock):
    def forward(self, x, index, *args, **kwargs):
        matrix = self.remap_matrix

        n_batch, _, n_val = x.shape
        new_nodes, _ = matrix.shape

        x = x.permute(2, 0, 1).flatten()
        x_unpooled = torch.zeros([new_nodes, n_batch * n_val], dtype=x.dtype, device=x.device)
        row, col = index
        x_unpooled =  torch.index_put(x_unpooled, (row, col), x)
        x_unpooled = x_unpooled.reshape(new_nodes, n_val, n_batch).permute(2, 0, 1)
        return x_unpooled