import torch


def is_available() -> bool:
    return torch.cuda.is_available()


try:
    _dynamo_disable = torch._dynamo.disable
except Exception:

    def _dynamo_disable(func):
        return func


@_dynamo_disable
def hard_voxelize_legacy_sorted(points: torch.Tensor, voxels: torch.Tensor,
                                coors: torch.Tensor,
                                num_points_per_voxel: torch.Tensor,
                                voxel_size, coors_range, max_points: int,
                                max_voxels: int, ndim: int = 3) -> int:
    """Deterministic hard voxelization without the CPU duplicate loop."""
    if ndim != 3:
        raise ValueError('hard_voxelize_legacy_sorted only supports ndim=3')
    if points.numel() == 0:
        return 0

    vx, vy, vz = [float(v) for v in voxel_size]
    x_min, y_min, z_min, x_max, y_max, z_max = [
        float(v) for v in coors_range
    ]
    grid_x = round((x_max - x_min) / vx)
    grid_y = round((y_max - y_min) / vy)
    grid_z = round((z_max - z_min) / vz)

    cx = torch.floor((points[:, 0] - x_min) / vx).to(torch.int64)
    cy = torch.floor((points[:, 1] - y_min) / vy).to(torch.int64)
    cz = torch.floor((points[:, 2] - z_min) / vz).to(torch.int64)
    valid = ((cx >= 0) & (cx < grid_x) & (cy >= 0) & (cy < grid_y)
             & (cz >= 0) & (cz < grid_z))
    valid_idx = valid.nonzero(as_tuple=False).flatten()
    if valid_idx.numel() == 0:
        return 0

    keys = cx * (grid_y * grid_z) + cy * grid_z + cz
    keys_valid = keys[valid_idx]

    order = torch.argsort(keys_valid, stable=True)
    keys_sorted = keys_valid[order]
    idx_sorted = valid_idx[order]

    new_group = torch.empty_like(keys_sorted, dtype=torch.bool)
    new_group[0] = True
    if keys_sorted.numel() > 1:
        new_group[1:] = keys_sorted[1:] != keys_sorted[:-1]

    group_starts = new_group.nonzero(as_tuple=False).flatten()
    num_groups = int(group_starts.numel())
    if num_groups == 0:
        return 0

    group_ids = torch.cumsum(new_group.to(torch.int64), dim=0) - 1
    sorted_pos = torch.arange(keys_sorted.numel(),
                              device=points.device,
                              dtype=torch.int64)
    rank_sorted = sorted_pos - group_starts[group_ids]

    end = torch.tensor([keys_sorted.numel()],
                       device=points.device,
                       dtype=torch.int64)
    group_ends = torch.cat((group_starts[1:], end))
    group_counts = group_ends - group_starts

    first_idx_by_key = idx_sorted[group_starts]
    group_order = torch.argsort(first_idx_by_key, stable=True)

    voxel_id_by_group = torch.empty((num_groups, ),
                                    device=points.device,
                                    dtype=torch.int64)
    voxel_id_by_group[group_order] = torch.arange(num_groups,
                                                  device=points.device,
                                                  dtype=torch.int64)
    voxel_id_sorted = voxel_id_by_group[group_ids]

    keep = (voxel_id_sorted < max_voxels) & (rank_sorted < max_points)
    if keep.any():
        voxels[voxel_id_sorted[keep], rank_sorted[keep]] = points[
            idx_sorted[keep]]

    num_out = min(num_groups, int(max_voxels))
    selected_groups = group_order[:num_out]
    coors_by_key = torch.stack(
        (cx[idx_sorted[group_starts]], cy[idx_sorted[group_starts]],
         cz[idx_sorted[group_starts]]),
        dim=1)
    coors[:num_out] = coors_by_key[selected_groups].to(coors.dtype)
    counts = torch.minimum(
        group_counts[selected_groups],
        torch.full((num_out, ),
                   int(max_points),
                   device=points.device,
                   dtype=torch.int64))
    num_points_per_voxel[:num_out] = counts.to(num_points_per_voxel.dtype)
    return num_out
