import torch
from torch.autograd import Function

from pointops._C import knn_query_cuda, random_ball_query_cuda, ball_query_cuda


class KNNQuery(Function):
    @staticmethod
    def forward(ctx, nsample, xyz, offset, new_xyz=None, new_offset=None):
        """
        input: coords: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b)
        output: idx: (m, nsample) -1 is placeholder, dist2: (m, nsample)
        """
        if new_xyz is None or new_offset is None:
            new_xyz = xyz
            new_offset = offset
        assert xyz.is_contiguous() and new_xyz.is_contiguous()
        m = new_xyz.shape[0]
        idx = torch.cuda.IntTensor(m, nsample).zero_()
        dist2 = torch.cuda.FloatTensor(m, nsample).zero_()
        knn_query_cuda(
            m, nsample, xyz, new_xyz, offset.int(), new_offset.int(), idx, dist2
        )
        return idx, torch.sqrt(dist2)


class RandomBallQuery(Function):
    """Random Ball Query.

    Find nearby points in spherical space.
    """

    @staticmethod
    def forward(
        ctx, nsample, max_radius, min_radius, xyz, offset, new_xyz=None, new_offset=None
    ):
        """
        input: coords: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b)
        output: idx: (m, nsample), dist2: (m, nsample)
        """
        if new_xyz is None or new_offset is None:
            new_xyz = xyz
            new_offset = offset
        assert xyz.is_contiguous() and new_xyz.is_contiguous()
        assert min_radius < max_radius

        m = new_xyz.shape[0]
        order = []
        for k in range(offset.shape[0]):
            s_k, e_k = (0, offset[0]) if k == 0 else (offset[k - 1], offset[k])
            order.append(
                torch.randperm(e_k - s_k, dtype=torch.int32, device=offset.device) + s_k
            )
        order = torch.cat(order, dim=0)
        idx = torch.cuda.IntTensor(m, nsample).zero_()
        dist2 = torch.cuda.FloatTensor(m, nsample).zero_()
        random_ball_query_cuda(
            m,
            nsample,
            min_radius,
            max_radius,
            order,
            xyz,
            new_xyz,
            offset.int(),
            new_offset.int(),
            idx,
            dist2,
        )
        return idx, torch.sqrt(dist2)


class BallQuery(Function):
    """Ball Query.

    Find nearby points in spherical space.
    """

    @staticmethod
    def forward(
        ctx, nsample, max_radius, min_radius, xyz, offset, new_xyz=None, new_offset=None
    ):
        """
        input: coords: (n, 3), new_xyz: (m, 3), offset: (b), new_offset: (b)
        output: idx: (m, nsample), dist2: (m, nsample)
        """
        if new_xyz is None or new_offset is None:
            new_xyz = xyz
            new_offset = offset
        assert xyz.is_contiguous() and new_xyz.is_contiguous()
        assert min_radius < max_radius

        m = new_xyz.shape[0]
        idx = torch.cuda.IntTensor(m, nsample).zero_()
        dist2 = torch.cuda.FloatTensor(m, nsample).zero_()
        ball_query_cuda(
            m,
            nsample,
            min_radius,
            max_radius,
            xyz,
            new_xyz,
            offset.int(),
            new_offset.int(),
            idx,
            dist2,
        )
        return idx, torch.sqrt(dist2)


knn_query = KNNQuery.apply
ball_query = BallQuery.apply
random_ball_query = RandomBallQuery.apply