import torch


def seq_continuous_group(xyz,sample_idx,nsample):
    """
    Input:
        nsample: max sample number in local region
        sample_idx: the index generated by previous sample step
        xyz: all points, [B, N, 3]
    Return:
        group_idx: grouped points index, [B, npoint, nsample]
    """
    B,N,_=xyz.shape
    _,npoint=sample_idx.shape
    idx=torch.zeros((B,npoint,nsample),dtype=torch.long,device=xyz.device)
    for i in range(B):
        for j in range(npoint):
            #check range
            start=sample_idx[i,j]
            if (start+nsample)<N:
                idx[i,j]=torch.arange(start,start+nsample)
            else:
                # num_overflow=start+nsample-N+1
                # num_in=N-num_overflow
                # idx[i,j][:num_in]=torch.arange(start,start+num_in)
                # idx[i,j][num_in:]=N-1
                idx[i,j]=torch.arange(N-nsample,N)
    tmp=idx[:,:,nsample//2]
    idx[:,:,nsample//2]=idx[:,:,0]
    idx[:,:,0]=tmp
    return idx