from typing import Tuple

import torch
from torch.autograd import Function
import torch.nn as nn

from metrics.pointops import pointops_cuda

import numpy as np


class FurthestSampling(Function):
    @staticmethod
    def forward(ctx, xyz, m):
        """
        input: xyz: (b, n, 3) and n > m, m: int32
        output: idx: (b, m)
        """
        assert xyz.is_contiguous()
        b, n, _ = xyz.size()
        idx = torch.cuda.IntTensor(b, m)
        temp = torch.cuda.FloatTensor(b, n).fill_(1e10)
        pointops_cuda.furthestsampling_cuda(b, n, m, xyz, temp, idx)
        return idx

    @staticmethod
    def backward(xyz, a=None):
        return None, None

furthestsampling = FurthestSampling.apply


class Gathering(Function):
    @staticmethod
    def forward(ctx, features, idx):
        """
        input: features: (b, c, n), idx : (b, m) tensor
        output: (b, c, m)
        """
        assert features.is_contiguous()
        assert idx.is_contiguous()
        b, c, n = features.size()
        m = idx.size(1)
        output = torch.cuda.FloatTensor(b, c, m)
        pointops_cuda.gathering_forward_cuda(b, c, n, m, features, idx, output)
        ctx.for_backwards = (idx, c, n)
        return output

    @staticmethod
    def backward(ctx, grad_out):
        idx, c, n = ctx.for_backwards
        b, m = idx.size()
        grad_features = torch.cuda.FloatTensor(b, c, n).zero_()
        grad_out_data = grad_out.data.contiguous()
        pointops_cuda.gathering_backward_cuda(b, c, n, m, grad_out_data, idx, grad_features.data)
        return grad_features, None

gathering = Gathering.apply


class NearestNeighbor(Function):
    @staticmethod
    def forward(ctx, unknown: torch.Tensor, known: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Find the three nearest neighbors of unknown in known
        input: unknown: (b, n, 3), known: (b, m, 3)
        output: dist2: (b, n, 3) l2 distance to the three nearest neighbors
                idx: (b, n, 3) index of 3 nearest neighbors
        """
        assert unknown.is_contiguous()
        assert known.is_contiguous()
        b, n, _ = unknown.size()
        m = known.size(1)
        dist2 = torch.cuda.FloatTensor(b, n, 3)
        idx = torch.cuda.IntTensor(b, n, 3)
        pointops_cuda.nearestneighbor_cuda(b, n, m, unknown, known, dist2, idx)
        return torch.sqrt(dist2), idx

    @staticmethod
    def backward(ctx, a=None, b=None):
        return None, None

nearestneighbor = NearestNeighbor.apply


class Interpolation(Function):
    @staticmethod
    def forward(ctx, features: torch.Tensor, idx: torch.Tensor, weight: torch.Tensor) -> torch.Tensor:
        """
        Performs weight linear interpolation on 3 features
        input: features: (b, c, m) features descriptors to be interpolated from
               idx: (b, n, 3) three nearest neighbors of the target features in features
               weight: (b, n, 3) weights
        output: (b, c, n) tensor of the interpolated features
        """
        assert features.is_contiguous()
        assert idx.is_contiguous()
        assert weight.is_contiguous()
        b, c, m = features.size()
        n = idx.size(1)
        ctx.interpolation_for_backward = (idx, weight, m)
        output = torch.cuda.FloatTensor(b, c, n)
        pointops_cuda.interpolation_forward_cuda(b, c, m, n, features, idx, weight, output)
        return output

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        input: grad_out: (b, c, n)
        output: grad_features: (b, c, m), None, None
        """
        idx, weight, m = ctx.interpolation_for_backward
        b, c, n = grad_out.size()
        grad_features = torch.cuda.FloatTensor(b, c, m).zero_()
        grad_out_data = grad_out.data.contiguous()
        pointops_cuda.interpolation_backward_cuda(b, c, n, m, grad_out_data, idx, weight, grad_features.data)
        return grad_features, None, None

interpolation = Interpolation.apply


class Grouping(Function):
    @staticmethod
    def forward(ctx, features: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
        """
        input: features: (b, c, n), idx : (b, m, nsample) containing the indicies of features to group with
        output: (b, c, m, nsample)
        """
        assert features.is_contiguous()
        assert idx.is_contiguous()
        b, c, n = features.size()
        _, m, nsample = idx.size()
        output = torch.cuda.FloatTensor(b, c, m, nsample)
        pointops_cuda.grouping_forward_cuda(b, c, n, m, nsample, features, idx, output)
        ctx.for_backwards = (idx, n)
        return output

    @staticmethod
    def backward(ctx, grad_out: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        input: grad_out: (b, c, m, nsample)
        output: (b, c, n), None
        """
        idx, n = ctx.for_backwards
        b, c, m, nsample = grad_out.size()
        grad_features = torch.cuda.FloatTensor(b, c, n).zero_()
        grad_out_data = grad_out.data.contiguous()
        pointops_cuda.grouping_backward_cuda(b, c, n, m, nsample, grad_out_data, idx, grad_features.data)
        return grad_features, None

grouping = Grouping.apply


class GroupingInt(Function):
    @staticmethod
    def forward(ctx, features: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
        """
        input: features: (b, c, n), idx : (b, m, nsample) containing the indicies of features to group with
        output: (b, c, m, nsample)
        """
        assert features.is_contiguous()
        assert idx.is_contiguous()
        b, c, n = features.size()
        _, m, nsample = idx.size()
        output = torch.cuda.LongTensor(b, c, m, nsample)
        pointops_cuda.grouping_int_forward_cuda(b, c, n, m, nsample, features, idx, output)
        return output

    @staticmethod
    def backward(ctx, a=None):
        return None, None

grouping_int = GroupingInt.apply


class BallQuery(Function):
    @staticmethod
    def forward(ctx, radius: float, nsample: int, xyz: torch.Tensor, new_xyz: torch.Tensor) -> torch.Tensor:
        """
        input: radius: float, radius of the balls
               nsample: int, maximum number of features in the balls
               xyz: torch.Tensor, (b, n, 3) xyz coordinates of the features
               new_xyz: torch.Tensor, (b, m, 3) centers of the ball query
        output: (b, m, nsample) tensor with the indicies of the features that form the query balls
        """
        assert xyz.is_contiguous()
        assert new_xyz.is_contiguous()
        b, n, _ = xyz.size()
        m = new_xyz.size(1)
        idx = torch.cuda.IntTensor(b, m, nsample).zero_()
        pointops_cuda.ballquery_cuda(b, n, m, radius, nsample, new_xyz, xyz, idx)
        return idx

    @staticmethod
    def backward(ctx, a=None):
        return None, None, None, None

ballquery = BallQuery.apply


class FeatureDistribute(Function):
    @staticmethod
    def forward(ctx, max_xyz: torch.Tensor, xyz: torch.Tensor) -> torch.Tensor:
        """
        :param ctx:
        :param max_xyz: (b, n, 3)
        :param xyz: (b, m, 3)
        :return: distribute_idx: (b, m)
        """
        assert max_xyz.is_contiguous()
        assert xyz.is_contiguous()
        b, n, _ = max_xyz.size()
        m = xyz.size(1)
        distribute_idx = torch.cuda.IntTensor(b, m).zero_()
        pointops_cuda.featuredistribute_cuda(b, n, m, max_xyz, xyz, distribute_idx)
        return distribute_idx

    @staticmethod
    def backward(ctx, a=None):
        return None, None

featuredistribute = FeatureDistribute.apply


class FeatureGather(Function):
    @staticmethod
    def forward(ctx, max_feature: torch.Tensor, distribute_idx: torch.Tensor) -> torch.Tensor:
        '''
        :param ctx:
        :param max_feature: (b, c, n)
        :param distribute_idx: (b, m)
        :return: distribute_feature: (b, c, m)
        '''
        assert max_feature.is_contiguous()
        assert distribute_idx.is_contiguous()
        b, c, n = max_feature.size()
        m = distribute_idx.size(1)
        distribute_feature = torch.cuda.FloatTensor(b, c, m).zero_()
        pointops_cuda.featuregather_forward_cuda(b, n, m, c, max_feature, distribute_idx, distribute_feature)
        ctx.for_backwards = (distribute_idx, n)
        return distribute_feature

    @staticmethod
    def backward(ctx, grad_distribute_feature: torch.Tensor):
        '''
        :param ctx:
        :param grad_distribute_feature: (b, c, m)
        :return: grad_max_feature: (b, c, n),    None
        '''
        distribute_idx, n = ctx.for_backwards
        b, c, m = grad_distribute_feature.size()
        grad_max_feature = torch.cuda.FloatTensor(b, c, n).zero_()
        grad_distribute_feature_data = grad_distribute_feature.data.contiguous()
        pointops_cuda.featuregather_backward_cuda(b, n, m, c, grad_distribute_feature_data, distribute_idx, grad_max_feature.data)
        return grad_max_feature, None

featuregather = FeatureGather.apply


class LabelStatBallRange(Function):
    @staticmethod
    def forward(ctx, radius: float, xyz: torch.Tensor, new_xyz: torch.Tensor, label_stat: torch.Tensor) -> torch.Tensor:
        '''
        :param ctx:
        :param radius:
        :param xyz: (b, n, 3)
        :param new_xyz: (b, m, 3)
        :param label_stat: (b, n, nclass)
        :return: new_label_stat: (b, m, nclass)
        '''
        assert xyz.is_contiguous()
        assert new_xyz.is_contiguous()
        assert label_stat.is_contiguous()

        b, n, nclass = label_stat.size()
        m = new_xyz.size(1)
        new_label_stat = torch.cuda.IntTensor(b, m, nclass).zero_()
        pointops_cuda.labelstat_ballrange_cuda(b, n, m, radius, nclass, new_xyz, xyz, label_stat, new_label_stat)

        return new_label_stat

    @staticmethod
    def backward(ctx, a=None):
        return None, None, None, None

labelstat_ballrange = LabelStatBallRange.apply


class LabelStatIdx(Function):
    @staticmethod
    def forward(ctx, nsample: int, label_stat: torch.Tensor, idx: torch.Tensor) -> torch.Tensor:
        '''
        :param ctx:
        :param nsample:
        :param label_stat: (b, n, nclass)
        :param idx: (b, m, nsample)
        :return: new_label_stat: (b, m, nclass)
        '''
        assert label_stat.is_contiguous()
        assert idx.is_contiguous()

        b, n, nclass = label_stat.size()
        m = idx.size(1)
        new_label_stat = torch.cuda.IntTensor(b, m, nclass).zero_()
        pointops_cuda.labelstat_idx_cuda(b, n, m, nsample, nclass, label_stat, idx, new_label_stat)

        return new_label_stat

    @staticmethod
    def backward(ctx, a=None):
        return None, None, None

labelstat_idx = LabelStatIdx.apply


class LabelStatAndBallQuery(Function):
    @staticmethod
    def forward(ctx, radius: float, nsample: int, xyz: torch.Tensor, new_xyz: torch.Tensor, label_stat: torch.Tensor):
        '''
        :param ctx:
        :param radius:
        :param nsample:
        :param xyz: (b, n, 3)
        :param new_xyz: (b, m, 3)
        :param label_stat: (b, n, nclass)
        :return: new_label_stat: (b, m, nclass)  idx: (b, m, nsample)
        '''
        assert xyz.is_contiguous()
        assert new_xyz.is_contiguous()
        assert label_stat.is_contiguous()

        b, n, nclass = label_stat.size()
        m = new_xyz.size(1)
        new_label_stat = torch.cuda.IntTensor(b, m, nclass).zero_()
        idx = torch.cuda.IntTensor(b, m, nsample).zero_()

        pointops_cuda.labelstat_and_ballquery_cuda(b, n, m, radius, nsample, nclass, new_xyz, xyz, label_stat, idx, new_label_stat)

        return new_label_stat, idx

    @staticmethod
    def backward(ctx, a=None, b=None):
        return None, None, None, None, None

labelstat_and_ballquery = LabelStatAndBallQuery.apply


def pairwise_distances(x, y=None):
    '''
    Input: x is a Nxd matrix
           y is an optional Mxd matirx
    Output: dist is a NxM matrix where dist[i,j] is the square norm between x[i,:] and y[j,:]
            if y is not given then use 'y=x'.
    i.e. dist[i,j] = ||x[i,:]-y[j,:]||^2
    '''
    x_norm = (x ** 2).sum(1).view(-1, 1)
    if y is not None:
        y_t = torch.transpose(y, 0, 1)
        y_norm = (y ** 2).sum(1).view(1, -1)
    else:
        y_t = torch.transpose(x, 0, 1)
        y_norm = x_norm.view(1, -1)
    dist = x_norm + y_norm - 2.0 * torch.mm(x, y_t)
    import numpy as np
    return torch.clamp(dist, 0.0, np.inf)


class KNNQueryNaive(Function):
    @staticmethod
    def forward(ctx, nsample: int, xyz: torch.Tensor, new_xyz: torch.Tensor = None) -> Tuple[torch.Tensor]:
        """
        KNN Indexing
        input: nsample: int32, Number of neighbor
               xyz: (b, n, 3) coordinates of the features
               new_xyz: (b, m, 3) centriods
            output: idx: (b, m, nsample)
        """
        if new_xyz is None:
            new_xyz = xyz
        b, m, _ = new_xyz.size()
        n = xyz.size(1)

        '''
        idx = torch.zeros(b, m, nsample).int().cuda()
        for i in range(b):
            dist = pairwise_distances(new_xyz[i, :, :], xyz[i, :, :])
            [_, idxs] = torch.sort(dist, dim=1)
            idx[i, :, :] = idxs[:, 0:nsample]
        '''

        # '''
        # new_xyz_repeat = new_xyz.repeat(1, 1, n).view(b, m * n, 3)
        # xyz_repeat = xyz.repeat(1, m, 1).view(b, m * n, 3)
        # dist = (new_xyz_repeat - xyz_repeat).pow(2).sum(dim=2).view(b, m, n)
        dist = (new_xyz.repeat(1, 1, n).view(b, m * n, 3) - xyz.repeat(1, m, 1).view(b, m * n, 3)).pow(2).sum(dim=2).view(b, m, n)
        [_, idxs] = torch.sort(dist, dim=2)
        idx = idxs[:, :, 0:nsample].int()
        # '''
        return idx

    @staticmethod
    def backward(ctx):
        return None, None, None

knnquery_naive = KNNQueryNaive.apply


class KNNQuery(Function):
    @staticmethod
    def forward(ctx, nsample: int, xyz: torch.Tensor, new_xyz: torch.Tensor = None) -> Tuple[torch.Tensor]:
        """
        KNN Indexing
        input: nsample: int32, Number of neighbor
               xyz: (b, n, 3) coordinates of the features
               new_xyz: (b, m, 3) centriods
            output: idx: (b, m, nsample)
                   ( dist2: (b, m, nsample) )
        """
        if new_xyz is None:
            new_xyz = xyz
        assert xyz.is_contiguous()
        assert new_xyz.is_contiguous()
        b, m, _ = new_xyz.size()
        n = xyz.size(1)
        idx = torch.cuda.IntTensor(b, m, nsample).zero_()
        dist2 = torch.cuda.FloatTensor(b, m, nsample).zero_()
        pointops_cuda.knnquery_cuda(b, n, m, nsample, xyz, new_xyz, idx, dist2)
        return idx

    @staticmethod
    def backward(ctx, a=None):
        return None, None, None

knnquery = KNNQuery.apply


class KNNQueryExclude(Function):
    @staticmethod
    def forward(ctx, nsample: int, xyz: torch.Tensor, new_xyz: torch.Tensor = None) -> Tuple[torch.Tensor]:
        """
        KNN Indexing
        input: nsample: int32, Number of neighbor
               xyz: (b, n, 3) coordinates of the features
               new_xyz: (b, m, 3) centriods
            output: new_features: (b, m, nsample)
        """
        if new_xyz is None:
            new_xyz = xyz
        b, m, _ = new_xyz.size()
        n = xyz.size(1)

        '''
        idx = torch.zeros(b, m, nsample).int().cuda()
        for i in range(b):
            dist = pairwise_distances(new_xyz[i, :, :], xyz[i, :, :])
            [_, idxs] = torch.sort(dist, dim=1)
            idx[i, :, :] = idxs[:, 0:nsample]
        '''

        # '''
        # new_xyz_repeat = new_xyz.repeat(1, 1, n).view(b, m * n, 3)
        # xyz_repeat = xyz.repeat(1, m, 1).view(b, m * n, 3)
        # dist = (new_xyz_repeat - xyz_repeat).pow(2).sum(dim=2).view(b, m, n)
        dist = (new_xyz.repeat(1, 1, n).view(b, m * n, 3) - xyz.repeat(1, m, 1).view(b, m * n, 3)).pow(2).sum(dim=2).view(b, m, n)
        [_, idxs] = torch.sort(dist, dim=2)
        idx = idxs[:, :, 1:nsample+1].int()
        # '''
        return idx

    @staticmethod
    def backward(ctx):
        return None, None, None

knnquery_exclude = KNNQueryExclude.apply

class Le_QueryAndGroup_SameSize(nn.Module):
    """
    Groups with a ball query of radius
    parameters:
        radius: float32, Radius of ball
        nsample: int32, Maximum number of features to gather in the ball
    """
    def __init__(self, radius=None, nsample=32, use_xyz=True):
        super(Le_QueryAndGroup_SameSize, self).__init__()
        self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz

    def forward(self, xyz: torch.Tensor, new_xyz: torch.Tensor = None, features: torch.Tensor = None, idx: torch.Tensor = None) -> torch.Tensor:
        """
        input: xyz: (b, n, 3) coordinates of the features
               new_xyz: (b, n, 3) centriods
               features: (b, c, n)
               idx: idx of neighbors
               # idxs: (b, n)
        output: new_features: (b, c+3, m, nsample)
              #  grouped_idxs: (b, m, nsample)
        """
        assert xyz.size() == new_xyz.size()
        if new_xyz is None:
            new_xyz = xyz
        if idx is None:
            if self.radius is not None:
                idx = ballquery(self.radius, self.nsample, xyz, new_xyz)
            else:
                # idx = knnquery_naive(self.nsample, xyz, new_xyz)   # (b, m, nsample)
                idx = knnquery(self.nsample, xyz, new_xyz)  # (b, m, nsample)
        xyz_trans = xyz.transpose(1, 2).contiguous()
        
        grouped_xyz = grouping(xyz_trans, idx)  # (b, 3, m, nsample)
        # grouped_idxs = grouping(idxs.unsqueeze(1).float(), idx).squeeze(1).int()  # (b, m, nsample)
        grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1)

        if features is not None:
            grouped_features = grouping(features, idx)  # (b, c, m, nsample)

            if self.use_xyz:
                #new_features = torch.cat([grouped_xyz, grouped_features], dim=1)  # (b, c+3, m, nsample)   # le
                new_features = grouped_features  # (b, c, m, nsample)
            else:
                new_features = grouped_features
        else:
            assert self.use_xyz, "Cannot have not features and not use xyz as a feature!"
            new_features = grouped_xyz
        return grouped_xyz, new_features


class QueryAndGroup(nn.Module):
    """
    Groups with a ball query of radius
    parameters:
        radius: float32, Radius of ball
        nsample: int32, Maximum number of features to gather in the ball
    """
    def __init__(self, radius=None, nsample=32, use_xyz=True):
        super(QueryAndGroup, self).__init__()
        self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz

    def forward(self, xyz: torch.Tensor, new_xyz: torch.Tensor = None, features: torch.Tensor = None, idx: torch.Tensor = None) -> torch.Tensor:
        """
        input: xyz: (b, n, 3) coordinates of the features
               new_xyz: (b, m, 3) centriods
               features: (b, c, n)
               idx: idx of neighbors
               # idxs: (b, n)
        output: new_features: (b, c+3, m, nsample)
              #  grouped_idxs: (b, m, nsample)
        """
        if new_xyz is None:
            new_xyz = xyz
        if idx is None:
            if self.radius is not None:
                idx = ballquery(self.radius, self.nsample, xyz, new_xyz)
            else:
                # idx = knnquery_naive(self.nsample, xyz, new_xyz)   # (b, m, nsample)
                idx = knnquery(self.nsample, xyz, new_xyz)  # (b, m, nsample)
        xyz_trans = xyz.transpose(1, 2).contiguous()
        grouped_xyz = grouping(xyz_trans, idx)  # (b, 3, m, nsample)
        # grouped_idxs = grouping(idxs.unsqueeze(1).float(), idx).squeeze(1).int()  # (b, m, nsample)

        grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1)
        if features is not None:
            grouped_features = grouping(features, idx)
            if self.use_xyz:
                new_features = torch.cat([grouped_xyz, grouped_features], dim=1)  # (b, c+3, m, nsample)
            else:
                new_features = grouped_features
        else:
            assert self.use_xyz, "Cannot have not features and not use xyz as a feature!"
            new_features = grouped_xyz
        return new_features

class QueryAndGroup_Dilate(nn.Module):
    """
    Groups with a ball query of radius
    parameters:
        radius: float32, Radius of ball
        nsample: int32, Maximum number of features to gather in the ball
    """
    def __init__(self, radius=None, nsample=32, use_xyz=True):
        super(QueryAndGroup_Dilate, self).__init__()
        self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz

    def forward(self, xyz: torch.Tensor, new_xyz: torch.Tensor = None, features: torch.Tensor = None, idx: torch.Tensor = None) -> torch.Tensor:
        """
        input: xyz: (b, n, 3) coordinates of the features
               new_xyz: (b, m, 3) centriods
               features: (b, c, n)
               idx: idx of neighbors
               # idxs: (b, n)
        output: new_features: (b, c+3, m, nsample)
              #  grouped_idxs: (b, m, nsample)
        """
        if new_xyz is None:
            new_xyz = xyz
        if idx is None:
            if self.radius is not None:
                idx = ballquery(self.radius, 2*self.nsample, xyz, new_xyz)
            else:
                # idx = knnquery_naive(self.nsample, xyz, new_xyz)   # (b, m, nsample)
                idx = knnquery(2*self.nsample, xyz, new_xyz)  # (b, m, nsample)
        
        idx2 = np.array([i for i in range(2*self.nsample)])
        np.random.shuffle(idx2)
        idx2 = idx2[:self.nsample] 
        idx = idx[:, :, idx2]
        
        xyz_trans = xyz.transpose(1, 2).contiguous()
        grouped_xyz = grouping(xyz_trans, idx)  # (b, 3, m, nsample)
        # grouped_idxs = grouping(idxs.unsqueeze(1).float(), idx).squeeze(1).int()  # (b, m, nsample)

        grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1)
        if features is not None:
            grouped_features = grouping(features, idx)
            if self.use_xyz:
                new_features = torch.cat([grouped_xyz, grouped_features], dim=1)  # (b, c+3, m, nsample)
            else:
                new_features = grouped_features
        else:
            assert self.use_xyz, "Cannot have not features and not use xyz as a feature!"
            new_features = grouped_xyz
        return new_features

class Le_QueryAndGroup(nn.Module):
    """
    Groups with a ball query of radius
    parameters:
        radius: float32, Radius of ball
        nsample: int32, Maximum number of features to gather in the ball
    """
    def __init__(self, radius=None, nsample=32, use_xyz=True):
        super(Le_QueryAndGroup, self).__init__()
        self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz

    def forward(self, xyz: torch.Tensor, new_xyz: torch.Tensor = None, features: torch.Tensor = None, idx: torch.Tensor = None) -> torch.Tensor:
        """
        input: xyz: (b, n, 3) coordinates of the features
               new_xyz: (b, m, 3) centriods
               features: (b, c, n)
               idx: idx of neighbors
               # idxs: (b, n)
        output: new_features: (b, c+3, m, nsample)
              #  grouped_idxs: (b, m, nsample)
        """
        if new_xyz is None:
            new_xyz = xyz
        if idx is None:
            if self.radius is not None:
                idx = ballquery(self.radius, self.nsample, xyz, new_xyz)
            else:
                # idx = knnquery_naive(self.nsample, xyz, new_xyz)   # (b, m, nsample)
                idx = knnquery(self.nsample, xyz, new_xyz)  # (b, m, nsample)
        xyz_trans = xyz.transpose(1, 2).contiguous()
        
        grouped_xyz = grouping(xyz_trans, idx)  # (b, 3, m, nsample)
        # grouped_idxs = grouping(idxs.unsqueeze(1).float(), idx).squeeze(1).int()  # (b, m, nsample)
        grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1)

        if features is not None:
            grouped_features = grouping(features, idx)  # (b, c, m, nsample)

            if self.use_xyz:
                #new_features = torch.cat([grouped_xyz, grouped_features], dim=1)  # (b, c+3, m, nsample)   # le
                new_features = grouped_features  # (b, c, m, nsample)
            else:
                new_features = grouped_features
        else:
            assert self.use_xyz, "Cannot have not features and not use xyz as a feature!"
            new_features = grouped_xyz
        return grouped_xyz, new_features

class Gen_QueryAndGroupXYZ(nn.Module):
    """
    Groups with a ball query of radius
    parameters:
        radius: float32, Radius of ball
        nsample: int32, Maximum number of features to gather in the ball
    """
    def __init__(self, radius=None, nsample=32, use_xyz=True):
        super(Gen_QueryAndGroupXYZ, self).__init__()
        self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz

    #def forward(self, xyz: torch.Tensor, new_xyz: torch.Tensor = None, features: torch.Tensor = None, idx: torch.Tensor = None) -> torch.Tensor:
    def forward(self, xyz: torch.Tensor, new_xyz: torch.Tensor = None) -> torch.Tensor:
        """
        input: xyz: (b, n, 3) coordinates of the features
               new_xyz: (b, m, 3) centriods
               features: (b, c, n)
               idx: idx of neighbors
               # idxs: (b, n)
        output: new_features: (b, c+3, m, nsample)
              #  grouped_idxs: (b, m, nsample)
        """
        if new_xyz is None:
            new_xyz = xyz
        #if idx is None:
        if self.radius is not None:
            idx = ballquery(self.radius, self.nsample, xyz, new_xyz)
        else:
            idx = knnquery(self.nsample, xyz, new_xyz)  # (b, m, nsample)
        xyz_trans = xyz.transpose(1, 2).contiguous()    # BxNx3 -> Bx3xN
        
        grouped_xyz = grouping(xyz_trans, idx)  # (b, 3, m, nsample)

        return grouped_xyz

class Le_QueryAndGroup_OnlyFeature(nn.Module):
    """
    Groups with a ball query of radius
    parameters:
        radius: float32, Radius of ball
        nsample: int32, Maximum number of features to gather in the ball
    """
    def __init__(self, radius=None, nsample=32, use_xyz=True):
        super(Le_QueryAndGroup_OnlyFeature, self).__init__()
        self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz

    def forward(self, xyz: torch.Tensor, new_xyz: torch.Tensor = None, features: torch.Tensor = None, idx: torch.Tensor = None) -> torch.Tensor:
        """
        input: xyz: (b, n, 3) coordinates of the features
               new_xyz: (b, m, 3) centriods
               features: (b, c, n)
               idx: idx of neighbors
               # idxs: (b, n)
        output: new_features: (b, c+3, m, nsample)
              #  grouped_idxs: (b, m, nsample)
        """
        if new_xyz is None:
            new_xyz = xyz
        if idx is None:
            if self.radius is not None:
                idx = ballquery(self.radius, self.nsample, xyz, new_xyz)
            else:
                # idx = knnquery_naive(self.nsample, xyz, new_xyz)   # (b, m, nsample)
                idx = knnquery(self.nsample, xyz, new_xyz)  # (b, m, nsample)
        #xyz_trans = xyz.transpose(1, 2).contiguous()
        
        #grouped_xyz = grouping(xyz_trans, idx)  # (b, 3, m, nsample)
        # grouped_idxs = grouping(idxs.unsqueeze(1).float(), idx).squeeze(1).int()  # (b, m, nsample)
        #grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1)

        if features is not None:
            grouped_features = grouping(features, idx)  # (b, c, m, nsample)

            if self.use_xyz:
                #new_features = torch.cat([grouped_xyz, grouped_features], dim=1)  # (b, c+3, m, nsample)   # le
                new_features = grouped_features  # (b, c, m, nsample)
            else:
                new_features = grouped_features
        else:
            assert self.use_xyz, "Cannot have not features and not use xyz as a feature!"
            new_features = grouped_xyz
        return new_features

class GroupAll(nn.Module):
    """
    Groups all features
    """
    def __init__(self, use_xyz: bool = True):
        super(GroupAll, self).__init__()
        self.use_xyz = use_xyz

    def forward(self, xyz: torch.Tensor, new_xyz: torch.Tensor, features: torch.Tensor = None) -> Tuple[torch.Tensor]:
        """
        input: xyz: (b, n, 3) coordinates of the features
               new_xyz: ignored torch
               features: (b, c, n) descriptors of the features
        output: new_features: (b, c+3, 1, N) tensor
        """
        grouped_xyz = xyz.transpose(1, 2).unsqueeze(2)
        if features is not None:
            grouped_features = features.unsqueeze(2)
            if self.use_xyz:
                new_features = torch.cat([grouped_xyz, grouped_features], dim=1)  # (b, c+3, 1, n)
            else:
                new_features = grouped_features
        else:
            new_features = grouped_xyz
        return new_features



[((365, 11, 365, 41), 'torch.clamp', 'torch.clamp', ({(365, 23, 365, 27): 'dist', (365, 29, 365, 32): '(0.0)', (365, 34, 365, 40): 'np.inf'}, {}), '(dist, 0.0, np.inf)', False, 'import torch\n'), ((21, 14, 21, 40), 'torch.cuda.IntTensor', 'torch.cuda.IntTensor', ({(21, 35, 21, 36): 'b', (21, 38, 21, 39): 'm'}, {}), '(b, m)', False, 'import torch\n'), ((23, 8, 23, 68), 'metrics.pointops.pointops_cuda.furthestsampling_cuda', 'pointops_cuda.furthestsampling_cuda', ({(23, 44, 23, 45): 'b', (23, 47, 23, 48): 'n', (23, 50, 23, 51): 'm', (23, 53, 23, 56): 'xyz', (23, 58, 23, 62): 'temp', (23, 64, 23, 67): 'idx'}, {}), '(b, n, m, xyz, temp, idx)', False, 'from metrics.pointops import pointops_cuda\n'), ((44, 17, 44, 48), 'torch.cuda.FloatTensor', 'torch.cuda.FloatTensor', ({(44, 40, 44, 41): 'b', (44, 43, 44, 44): 'c', (44, 46, 44, 47): 'm'}, {}), '(b, c, m)', False, 'import torch\n'), ((45, 8, 45, 79), 'metrics.pointops.pointops_cuda.gathering_forward_cuda', 'pointops_cuda.gathering_forward_cuda', ({(45, 45, 45, 46): 'b', (45, 48, 45, 49): 'c', (45, 51, 45, 52): 'n', (45, 54, 45, 55): 'm', (45, 57, 45, 65): 'features', (45, 67, 45, 70): 'idx', (45, 72, 45, 78): 'output'}, {}), '(b, c, n, m, features, idx, output)', False, 'from metrics.pointops import pointops_cuda\n'), ((55, 8, 55, 97), 'metrics.pointops.pointops_cuda.gathering_backward_cuda', 'pointops_cuda.gathering_backward_cuda', ({(55, 46, 55, 47): 'b', (55, 49, 55, 50): 'c', (55, 52, 55, 53): 'n', (55, 55, 55, 56): 'm', (55, 58, 55, 71): 'grad_out_data', (55, 73, 55, 76): 'idx', (55, 78, 55, 96): 'grad_features.data'}, {}), '(b, c, n, m, grad_out_data, idx,\n    grad_features.data)', False, 'from metrics.pointops import pointops_cuda\n'), ((74, 16, 74, 47), 'torch.cuda.FloatTensor', 'torch.cuda.FloatTensor', ({(74, 39, 74, 40): 'b', (74, 42, 74, 43): 'n', (74, 45, 74, 46): '3'}, {}), '(b, n, 3)', False, 'import torch\n'), ((75, 14, 75, 43), 'torch.cuda.IntTensor', 'torch.cuda.IntTensor', ({(75, 35, 75, 36): 'b', (75, 38, 75, 39): 'n', (75, 41, 75, 42): '3'}, {}), '(b, n, 3)', False, 'import torch\n'), ((76, 8, 76, 79), 'metrics.pointops.pointops_cuda.nearestneighbor_cuda', 'pointops_cuda.nearestneighbor_cuda', ({(76, 43, 76, 44): 'b', (76, 46, 76, 47): 'n', (76, 49, 76, 50): 'm', (76, 52, 76, 59): 'unknown', (76, 61, 76, 66): 'known', (76, 68, 76, 73): 'dist2', (76, 75, 76, 78): 'idx'}, {}), '(b, n, m, unknown, known, dist2, idx)', False, 'from metrics.pointops import pointops_cuda\n'), ((102, 17, 102, 48), 'torch.cuda.FloatTensor', 'torch.cuda.FloatTensor', ({(102, 40, 102, 41): 'b', (102, 43, 102, 44): 'c', (102, 46, 102, 47): 'n'}, {}), '(b, c, n)', False, 'import torch\n'), ((103, 8, 103, 91), 'metrics.pointops.pointops_cuda.interpolation_forward_cuda', 'pointops_cuda.interpolation_forward_cuda', ({(103, 49, 103, 50): 'b', (103, 52, 103, 53): 'c', (103, 55, 103, 56): 'm', (103, 58, 103, 59): 'n', (103, 61, 103, 69): 'features', (103, 71, 103, 74): 'idx', (103, 76, 103, 82): 'weight', (103, 84, 103, 90): 'output'}, {}), '(b, c, m, n, features, idx, weight,\n    output)', False, 'from metrics.pointops import pointops_cuda\n'), ((116, 8, 116, 109), 'metrics.pointops.pointops_cuda.interpolation_backward_cuda', 'pointops_cuda.interpolation_backward_cuda', ({(116, 50, 116, 51): 'b', (116, 53, 116, 54): 'c', (116, 56, 116, 57): 'n', (116, 59, 116, 60): 'm', (116, 62, 116, 75): 'grad_out_data', (116, 77, 116, 80): 'idx', (116, 82, 116, 88): 'weight', (116, 90, 116, 108): 'grad_features.data'}, {}), '(b, c, n, m, grad_out_data, idx,\n    weight, grad_features.data)', False, 'from metrics.pointops import pointops_cuda\n'), ((133, 17, 133, 57), 'torch.cuda.FloatTensor', 'torch.cuda.FloatTensor', ({(133, 40, 133, 41): 'b', (133, 43, 133, 44): 'c', (133, 46, 133, 47): 'm', (133, 49, 133, 56): 'nsample'}, {}), '(b, c, m, nsample)', False, 'import torch\n'), ((134, 8, 134, 87), 'metrics.pointops.pointops_cuda.grouping_forward_cuda', 'pointops_cuda.grouping_forward_cuda', ({(134, 44, 134, 45): 'b', (134, 47, 134, 48): 'c', (134, 50, 134, 51): 'n', (134, 53, 134, 54): 'm', (134, 56, 134, 63): 'nsample', (134, 65, 134, 73): 'features', (134, 75, 134, 78): 'idx', (134, 80, 134, 86): 'output'}, {}), '(b, c, n, m, nsample, features, idx, output)', False, 'from metrics.pointops import pointops_cuda\n'), ((148, 8, 148, 105), 'metrics.pointops.pointops_cuda.grouping_backward_cuda', 'pointops_cuda.grouping_backward_cuda', ({(148, 45, 148, 46): 'b', (148, 48, 148, 49): 'c', (148, 51, 148, 52): 'n', (148, 54, 148, 55): 'm', (148, 57, 148, 64): 'nsample', (148, 66, 148, 79): 'grad_out_data', (148, 81, 148, 84): 'idx', (148, 86, 148, 104): 'grad_features.data'}, {}), '(b, c, n, m, nsample, grad_out_data,\n    idx, grad_features.data)', False, 'from metrics.pointops import pointops_cuda\n'), ((165, 17, 165, 56), 'torch.cuda.LongTensor', 'torch.cuda.LongTensor', ({(165, 39, 165, 40): 'b', (165, 42, 165, 43): 'c', (165, 45, 165, 46): 'm', (165, 48, 165, 55): 'nsample'}, {}), '(b, c, m, nsample)', False, 'import torch\n'), ((166, 8, 166, 91), 'metrics.pointops.pointops_cuda.grouping_int_forward_cuda', 'pointops_cuda.grouping_int_forward_cuda', ({(166, 48, 166, 49): 'b', (166, 51, 166, 52): 'c', (166, 54, 166, 55): 'n', (166, 57, 166, 58): 'm', (166, 60, 166, 67): 'nsample', (166, 69, 166, 77): 'features', (166, 79, 166, 82): 'idx', (166, 84, 166, 90): 'output'}, {}), '(b, c, n, m, nsample, features, idx,\n    output)', False, 'from metrics.pointops import pointops_cuda\n'), ((191, 8, 191, 81), 'metrics.pointops.pointops_cuda.ballquery_cuda', 'pointops_cuda.ballquery_cuda', ({(191, 37, 191, 38): 'b', (191, 40, 191, 41): 'n', (191, 43, 191, 44): 'm', (191, 46, 191, 52): 'radius', (191, 54, 191, 61): 'nsample', (191, 63, 191, 70): 'new_xyz', (191, 72, 191, 75): 'xyz', (191, 77, 191, 80): 'idx'}, {}), '(b, n, m, radius, nsample, new_xyz, xyz, idx)', False, 'from metrics.pointops import pointops_cuda\n'), ((215, 8, 215, 83), 'metrics.pointops.pointops_cuda.featuredistribute_cuda', 'pointops_cuda.featuredistribute_cuda', ({(215, 45, 215, 46): 'b', (215, 48, 215, 49): 'n', (215, 51, 215, 52): 'm', (215, 54, 215, 61): 'max_xyz', (215, 63, 215, 66): 'xyz', (215, 68, 215, 82): 'distribute_idx'}, {}), '(b, n, m, max_xyz, xyz, distribute_idx)', False, 'from metrics.pointops import pointops_cuda\n'), ((239, 8, 239, 109), 'metrics.pointops.pointops_cuda.featuregather_forward_cuda', 'pointops_cuda.featuregather_forward_cuda', ({(239, 49, 239, 50): 'b', (239, 52, 239, 53): 'n', (239, 55, 239, 56): 'm', (239, 58, 239, 59): 'c', (239, 61, 239, 72): 'max_feature', (239, 74, 239, 88): 'distribute_idx', (239, 90, 239, 108): 'distribute_feature'}, {}), '(b, n, m, c, max_feature,\n    distribute_idx, distribute_feature)', False, 'from metrics.pointops import pointops_cuda\n'), ((254, 8, 254, 130), 'metrics.pointops.pointops_cuda.featuregather_backward_cuda', 'pointops_cuda.featuregather_backward_cuda', ({(254, 50, 254, 51): 'b', (254, 53, 254, 54): 'n', (254, 56, 254, 57): 'm', (254, 59, 254, 60): 'c', (254, 62, 254, 90): 'grad_distribute_feature_data', (254, 92, 254, 106): 'distribute_idx', (254, 108, 254, 129): 'grad_max_feature.data'}, {}), '(b, n, m, c,\n    grad_distribute_feature_data, distribute_idx, grad_max_feature.data)', False, 'from metrics.pointops import pointops_cuda\n'), ((278, 8, 278, 113), 'metrics.pointops.pointops_cuda.labelstat_ballrange_cuda', 'pointops_cuda.labelstat_ballrange_cuda', ({(278, 47, 278, 48): 'b', (278, 50, 278, 51): 'n', (278, 53, 278, 54): 'm', (278, 56, 278, 62): 'radius', (278, 64, 278, 70): 'nclass', (278, 72, 278, 79): 'new_xyz', (278, 81, 278, 84): 'xyz', (278, 86, 278, 96): 'label_stat', (278, 98, 278, 112): 'new_label_stat'}, {}), '(b, n, m, radius, nclass, new_xyz,\n    xyz, label_stat, new_label_stat)', False, 'from metrics.pointops import pointops_cuda\n'), ((305, 8, 305, 99), 'metrics.pointops.pointops_cuda.labelstat_idx_cuda', 'pointops_cuda.labelstat_idx_cuda', ({(305, 41, 305, 42): 'b', (305, 44, 305, 45): 'n', (305, 47, 305, 48): 'm', (305, 50, 305, 57): 'nsample', (305, 59, 305, 65): 'nclass', (305, 67, 305, 77): 'label_stat', (305, 79, 305, 82): 'idx', (305, 84, 305, 98): 'new_label_stat'}, {}), '(b, n, m, nsample, nclass, label_stat, idx,\n    new_label_stat)', False, 'from metrics.pointops import pointops_cuda\n'), ((337, 8, 337, 131), 'metrics.pointops.pointops_cuda.labelstat_and_ballquery_cuda', 'pointops_cuda.labelstat_and_ballquery_cuda', ({(337, 51, 337, 52): 'b', (337, 54, 337, 55): 'n', (337, 57, 337, 58): 'm', (337, 60, 337, 66): 'radius', (337, 68, 337, 75): 'nsample', (337, 77, 337, 83): 'nclass', (337, 85, 337, 92): 'new_xyz', (337, 94, 337, 97): 'xyz', (337, 99, 337, 109): 'label_stat', (337, 111, 337, 114): 'idx', (337, 116, 337, 130): 'new_label_stat'}, {}), '(b, n, m, radius, nsample, nclass,\n    new_xyz, xyz, label_stat, idx, new_label_stat)', False, 'from metrics.pointops import pointops_cuda\n'), ((358, 14, 358, 38), 'torch.transpose', 'torch.transpose', ({(358, 30, 358, 31): 'y', (358, 33, 358, 34): '0', (358, 36, 358, 37): '1'}, {}), '(y, 0, 1)', False, 'import torch\n'), ((361, 14, 361, 38), 'torch.transpose', 'torch.transpose', ({(361, 30, 361, 31): 'x', (361, 33, 361, 34): '0', (361, 36, 361, 37): '1'}, {}), '(x, 0, 1)', False, 'import torch\n'), ((396, 20, 396, 43), 'torch.sort', 'torch.sort', (), '', False, 'import torch\n'), ((427, 8, 427, 79), 'metrics.pointops.pointops_cuda.knnquery_cuda', 'pointops_cuda.knnquery_cuda', ({(427, 36, 427, 37): 'b', (427, 39, 427, 40): 'n', (427, 42, 427, 43): 'm', (427, 45, 427, 52): 'nsample', (427, 54, 427, 57): 'xyz', (427, 59, 427, 66): 'new_xyz', (427, 68, 427, 71): 'idx', (427, 73, 427, 78): 'dist2'}, {}), '(b, n, m, nsample, xyz, new_xyz, idx, dist2)', False, 'from metrics.pointops import pointops_cuda\n'), ((465, 20, 465, 43), 'torch.sort', 'torch.sort', (), '', False, 'import torch\n'), ((602, 8, 602, 31), 'numpy.random.shuffle', 'np.random.shuffle', ({(602, 26, 602, 30): 'idx2'}, {}), '(idx2)', True, 'import numpy as np\n'), ((77, 15, 77, 32), 'torch.sqrt', 'torch.sqrt', ({(77, 26, 77, 31): 'dist2'}, {}), '(dist2)', False, 'import torch\n'), ((363, 35, 363, 51), 'torch.mm', 'torch.mm', ({(363, 44, 363, 45): 'x', (363, 47, 363, 50): 'y_t'}, {}), '(x, y_t)', False, 'import torch\n'), ((22, 15, 22, 43), 'torch.cuda.FloatTensor', 'torch.cuda.FloatTensor', ({(22, 38, 22, 39): 'b', (22, 41, 22, 42): 'n'}, {}), '(b, n)', False, 'import torch\n'), ((53, 24, 53, 55), 'torch.cuda.FloatTensor', 'torch.cuda.FloatTensor', ({(53, 47, 53, 48): 'b', (53, 50, 53, 51): 'c', (53, 53, 53, 54): 'n'}, {}), '(b, c, n)', False, 'import torch\n'), ((114, 24, 114, 55), 'torch.cuda.FloatTensor', 'torch.cuda.FloatTensor', ({(114, 47, 114, 48): 'b', (114, 50, 114, 51): 'c', (114, 53, 114, 54): 'm'}, {}), '(b, c, m)', False, 'import torch\n'), ((146, 24, 146, 55), 'torch.cuda.FloatTensor', 'torch.cuda.FloatTensor', ({(146, 47, 146, 48): 'b', (146, 50, 146, 51): 'c', (146, 53, 146, 54): 'n'}, {}), '(b, c, n)', False, 'import torch\n'), ((190, 14, 190, 49), 'torch.cuda.IntTensor', 'torch.cuda.IntTensor', ({(190, 35, 190, 36): 'b', (190, 38, 190, 39): 'm', (190, 41, 190, 48): 'nsample'}, {}), '(b, m, nsample)', False, 'import torch\n'), ((214, 25, 214, 51), 'torch.cuda.IntTensor', 'torch.cuda.IntTensor', ({(214, 46, 214, 47): 'b', (214, 49, 214, 50): 'm'}, {}), '(b, m)', False, 'import torch\n'), ((238, 29, 238, 60), 'torch.cuda.FloatTensor', 'torch.cuda.FloatTensor', ({(238, 52, 238, 53): 'b', (238, 55, 238, 56): 'c', (238, 58, 238, 59): 'm'}, {}), '(b, c, m)', False, 'import torch\n'), ((252, 27, 252, 58), 'torch.cuda.FloatTensor', 'torch.cuda.FloatTensor', ({(252, 50, 252, 51): 'b', (252, 53, 252, 54): 'c', (252, 56, 252, 57): 'n'}, {}), '(b, c, n)', False, 'import torch\n'), ((277, 25, 277, 59), 'torch.cuda.IntTensor', 'torch.cuda.IntTensor', ({(277, 46, 277, 47): 'b', (277, 49, 277, 50): 'm', (277, 52, 277, 58): 'nclass'}, {}), '(b, m, nclass)', False, 'import torch\n'), ((304, 25, 304, 59), 'torch.cuda.IntTensor', 'torch.cuda.IntTensor', ({(304, 46, 304, 47): 'b', (304, 49, 304, 50): 'm', (304, 52, 304, 58): 'nclass'}, {}), '(b, m, nclass)', False, 'import torch\n'), ((334, 25, 334, 59), 'torch.cuda.IntTensor', 'torch.cuda.IntTensor', ({(334, 46, 334, 47): 'b', (334, 49, 334, 50): 'm', (334, 52, 334, 58): 'nclass'}, {}), '(b, m, nclass)', False, 'import torch\n'), ((335, 14, 335, 49), 'torch.cuda.IntTensor', 'torch.cuda.IntTensor', ({(335, 35, 335, 36): 'b', (335, 38, 335, 39): 'm', (335, 41, 335, 48): 'nsample'}, {}), '(b, m, nsample)', False, 'import torch\n'), ((425, 14, 425, 49), 'torch.cuda.IntTensor', 'torch.cuda.IntTensor', ({(425, 35, 425, 36): 'b', (425, 38, 425, 39): 'm', (425, 41, 425, 48): 'nsample'}, {}), '(b, m, nsample)', False, 'import torch\n'), ((426, 16, 426, 53), 'torch.cuda.FloatTensor', 'torch.cuda.FloatTensor', ({(426, 39, 426, 40): 'b', (426, 42, 426, 43): 'm', (426, 45, 426, 52): 'nsample'}, {}), '(b, m, nsample)', False, 'import torch\n'), ((563, 31, 563, 80), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((614, 31, 614, 80), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n'), ((772, 31, 772, 80), 'torch.cat', 'torch.cat', (), '', False, 'import torch\n')]