from collections import namedtuple

import torch


KNNResult = namedtuple("KNNResult", ["dists", "idx", "knn"])
BallQueryResult = namedtuple("BallQueryResult", ["dists", "idx", "knn"])


def knn_points(p1, p2, lengths1=None, lengths2=None, K=1, norm=2, return_nn=False, return_sorted=True):
    if p1.dim() != 3 or p2.dim() != 3:
        raise ValueError("p1 and p2 must have shape (N, P, D)")
    if p1.shape[0] != p2.shape[0] or p1.shape[2] != p2.shape[2]:
        raise ValueError("p1 and p2 must have matching batch and point dimensions")

    batch, p1_count, _ = p1.shape
    p2_count = p2.shape[1]
    k = min(K, p2_count)
    if k <= 0:
        empty_dists = p1.new_empty((batch, p1_count, 0))
        empty_idx = torch.empty((batch, p1_count, 0), dtype=torch.long, device=p1.device)
        empty_nn = p2.new_empty((batch, p1_count, 0, p2.shape[2])) if return_nn else None
        return KNNResult(empty_dists, empty_idx, empty_nn)

    dists = torch.cdist(p1.float(), p2.float(), p=norm)
    if norm == 2:
        dists = dists.square()

    if lengths2 is not None:
        arange = torch.arange(p2_count, device=p2.device).view(1, 1, p2_count)
        valid = arange < lengths2.to(device=p2.device).view(batch, 1, 1)
        dists = dists.masked_fill(~valid, torch.inf)

    dists_k, idx = torch.topk(dists, k=k, dim=-1, largest=False, sorted=return_sorted)

    if K > k:
        pad = K - k
        dists_k = torch.cat([dists_k, dists_k.new_full((batch, p1_count, pad), torch.inf)], dim=-1)
        idx = torch.cat([idx, idx.new_full((batch, p1_count, pad), -1)], dim=-1)

    if lengths1 is not None:
        arange = torch.arange(p1_count, device=p1.device).view(1, p1_count, 1)
        valid = arange < lengths1.to(device=p1.device).view(batch, 1, 1)
        dists_k = dists_k.masked_fill(~valid, torch.inf)
        idx = idx.masked_fill(~valid, -1)

    knn = None
    if return_nn:
        safe_idx = idx.clamp_min(0)
        gather_idx = safe_idx.unsqueeze(-1).expand(-1, -1, -1, p2.shape[2])
        points = p2.unsqueeze(1).expand(-1, p1_count, -1, -1)
        knn = torch.gather(points, 2, gather_idx)
        knn = knn.masked_fill(idx.unsqueeze(-1) < 0, 0)

    return KNNResult(dists_k.to(dtype=p1.dtype), idx, knn)


def ball_query(p1, p2, lengths1=None, lengths2=None, K=1, radius=0.2, return_nn=False):
    if p1.dim() != 3 or p2.dim() != 3:
        raise ValueError("p1 and p2 must have shape (N, P, D)")
    if p1.shape[0] != p2.shape[0] or p1.shape[2] != p2.shape[2]:
        raise ValueError("p1 and p2 must have matching batch and point dimensions")

    batch, p1_count, _ = p1.shape
    p2_count = p2.shape[1]
    k = max(int(K), 0)
    if k == 0:
        empty_dists = p1.new_empty((batch, p1_count, 0))
        empty_idx = torch.empty((batch, p1_count, 0), dtype=torch.long, device=p1.device)
        empty_nn = p2.new_empty((batch, p1_count, 0, p2.shape[2])) if return_nn else None
        return BallQueryResult(empty_dists, empty_idx, empty_nn)

    dists = torch.cdist(p1.float(), p2.float(), p=2).square()
    max_dist = float(radius) * float(radius)
    valid = dists <= max_dist

    if lengths2 is not None:
        arange = torch.arange(p2_count, device=p2.device).view(1, 1, p2_count)
        valid = valid & (arange < lengths2.to(device=p2.device).view(batch, 1, 1))

    masked = dists.masked_fill(~valid, torch.inf)
    take = min(k, p2_count)
    dists_k, idx = torch.topk(masked, k=take, dim=-1, largest=False, sorted=True)
    invalid = torch.isinf(dists_k)
    idx = idx.masked_fill(invalid, -1)
    dists_k = dists_k.masked_fill(invalid, 0)

    if k > take:
        pad = k - take
        dists_k = torch.cat([dists_k, dists_k.new_zeros((batch, p1_count, pad))], dim=-1)
        idx = torch.cat([idx, idx.new_full((batch, p1_count, pad), -1)], dim=-1)

    if lengths1 is not None:
        arange = torch.arange(p1_count, device=p1.device).view(1, p1_count, 1)
        valid_p1 = arange < lengths1.to(device=p1.device).view(batch, 1, 1)
        dists_k = dists_k.masked_fill(~valid_p1, 0)
        idx = idx.masked_fill(~valid_p1, -1)

    knn = None
    if return_nn:
        safe_idx = idx.clamp_min(0)
        gather_idx = safe_idx.unsqueeze(-1).expand(-1, -1, -1, p2.shape[2])
        points = p2.unsqueeze(1).expand(-1, p1_count, -1, -1)
        knn = torch.gather(points, 2, gather_idx)
        knn = knn.masked_fill(idx.unsqueeze(-1) < 0, 0)

    return BallQueryResult(dists_k.to(dtype=p1.dtype), idx, knn)