import numpy as np
from scipy.spatial.distance import cdist
from scipy.optimize import linear_sum_assignment


def match(
    set1: np.ndarray,
    set2: np.ndarray,
    xy_threshold=250,
    z_threshold=500,
):
    xy_sq_thresh, z_sq_thresh = xy_threshold**2, z_threshold**2
    numeric_inf = 10 * (xy_sq_thresh + z_sq_thresh)
    xy_sq_dist = cdist(set1[:, :2], set2[:, :2], metric="sqeuclidean")
    z_sq_dist = cdist(set1[:, 2:3], set2[:, 2:3], metric="sqeuclidean")
    valid_mask = (xy_sq_dist <= xy_sq_thresh) & (z_sq_dist <= z_sq_thresh)
    cost_matrix = np.where(valid_mask, xy_sq_dist + z_sq_dist, np.inf)
    cost_matrix_finite = np.where(valid_mask, cost_matrix, numeric_inf)
    matched_set1_idx, matched_set2_idx = linear_sum_assignment(cost_matrix_finite)
    valid = ~np.isinf(cost_matrix[matched_set1_idx, matched_set2_idx])
    return matched_set1_idx[valid], matched_set2_idx[valid]
