import numpy as np
import heapq
from dataclasses import dataclass
from typing import Optional


@dataclass
class KDNode:
    point: np.ndarray
    label: object
    axis: int
    left: Optional["KDNode"] = None
    right: Optional["KDNode"] = None


def build_kdtree(points, labels, depth=0):
    if len(points) == 0:
        return None
    axis = depth % points.shape[1]
    order = np.argsort(points[:, axis])
    points, labels = points[order], labels[order]
    mid = len(points) // 2
    return KDNode(
        point=points[mid],
        label=labels[mid],
        axis=axis,
        left=build_kdtree(points[:mid], labels[:mid], depth + 1),
        right=build_kdtree(points[mid + 1:], labels[mid + 1:], depth + 1),
    )


def knn_search(root, query, k):
    """Return the k nearest (dist, label) pairs via KD-tree pruning."""
    heap = []  # max-heap of (-dist, counter, label) to cap size at k
    counter = 0

    def visit(node):
        nonlocal counter
        if node is None:
            return
        dist = np.linalg.norm(node.point - query)
        if len(heap) < k:
            heapq.heappush(heap, (-dist, counter, node.label))
            counter += 1
        elif dist < -heap[0][0]:
            heapq.heapreplace(heap, (-dist, counter, node.label))
            counter += 1

        diff = query[node.axis] - node.point[node.axis]
        if diff < 0:
            near, far = node.left, node.right
        else:
            near, far = node.right, node.left
        visit(near)
        # only explore the far side if it could contain a closer point
        # than the current k-th best distance
        worst_best = -heap[0][0] if len(heap) == k else float("inf")
        if abs(diff) < worst_best:
            visit(far)

    visit(root)
    return sorted([(-d, lbl) for d, _, lbl in heap])


def knn_predict(root, query, k):
    neighbors = knn_search(root, query, k)
    labels = [lbl for _, lbl in neighbors]
    values, counts = np.unique(labels, return_counts=True)
    return values[np.argmax(counts)]
