input: Write a python function to implement a k-d tree for storing and querying k-nearest neighbors in a 2-Dimensional space.
  assert isinstance(kd_tree([(1, 2), (3, 4), (5, 6)]), KDTree)
  assert kd_tree([(1, 2), (3, 4), (5, 6)]).query((3, 3), k=1) == [(3, 4)]
  assert kd_tree([(1, 2), (3, 4), (5, 6)]).query((0, 0), k=2) == [(1, 2), (3, 4)]
output: from scipy.spatial import KDTree
  
  class KDTreeWrapper:
      def __init__(self, points):
          self.tree = KDTree(points)
          self.points = points
  
      def query(self, point, k=1):
          distances, indices = self.tree.query(point, k)
          if k == 1:
              return [self.points[indices]]
          else:
              return [self.points[i] for i in indices]
  
  def kd_tree(points):
      return KDTreeWrapper(points)