import torch
import torch.nn as nn

torch.backends.cuda.matmul.allow_tf32 = False

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.radius = 20.0
        self.nsample = 32

    def forward(self, xyz: torch.Tensor, new_xyz: torch.Tensor) -> torch.Tensor:
        """
        xyz: [B, N, 3] (Database points)
        new_xyz: [B, M, 3] (Query points)
        Output: [B, M, nsample] (Indices of neighbors)
        """

        dist2 = torch.sum((new_xyz.unsqueeze(2) - xyz.unsqueeze(1)) ** 2, dim=-1)

        r2 = self.radius ** 2
        mask = dist2 < r2
  
        _, indices = torch.sort(mask.int(), dim=-1, descending=True)

        idx_out = indices[:, :, :self.nsample]
    
        valid_count = mask.sum(dim=-1) # [B, M]
 
        range_vec = torch.arange(self.nsample, device=xyz.device).view(1, 1, -1)

        pad_mask = range_vec >= valid_count.unsqueeze(-1)

        first_indices = idx_out[:, :, 0:1] # [B, M, 1]
        
       
        idx_out = torch.where(pad_mask, first_indices, idx_out)
        
        
        return idx_out.int()

B = 4
N = 4096
M = 1024

def get_inputs():
    xyz = torch.randint(0, 100, (B, N, 3), device='cuda').float()
    new_xyz = torch.randint(0, 100, (B, M, 3), device='cuda').float()
    return [xyz, new_xyz]

def get_init_inputs():
    return []