def brute_force_knn(X_train, y_train, query, k):
    dists = np.linalg.norm(X_train - query, axis=1)
    nearest_idx = np.argsort(dists)[:k]
    labels = y_train[nearest_idx]
    values, counts = np.unique(labels, return_counts=True)
    return values[np.argmax(counts)]


if __name__ == "__main__":
    rng = np.random.default_rng(0)
    n = 400
    X = rng.normal(size=(n, 2))
    y = (X[:, 0] * X[:, 1] > 0).astype(int)  # XOR-like pattern

    split = int(0.8 * n)
    X_train, X_test = X[:split], X[split:]
    y_train, y_test = y[:split], y[split:]

    tree = build_kdtree(X_train, y_train)

    k = 5
    tree_preds = np.array(
        [knn_predict(tree, q, k) for q in X_test]
    )
    brute_preds = np.array(
        [brute_force_knn(X_train, y_train, q, k) for q in X_test]
    )

    agree = np.mean(tree_preds == brute_preds)
    accuracy = np.mean(tree_preds == y_test)
    print(f"KD-tree vs brute-force agreement: {agree:.3f}")
    print(f"KD-tree k-NN test accuracy: {accuracy:.3f}")
