import operator

import numpy as np


class KNN(object):

    def __init__(self, data_set, labels, k):
        self._N = data_set.shape[0]
        self._K = k
        self._labels = labels
        self._data_set = data_set
        pass

    def tell_label(self, vec):
        return tell_label(self._data_set, self._labels, self._K, vec)


def tell_label(data_set, labels, k, vec):
    N = data_set.shape[0]
    tiled_vecs = np.tile(vec, (N, 1))
    diff_mat = tiled_vecs - data_set
    diff_mat_square = diff_mat ** 2
    diff_sqare_sums = np.sum(diff_mat_square, axis=1)

    sorted_indexes = np.argsort(diff_sqare_sums)
    labels_count = {}
    for i in xrange(k):
        label = labels[sorted_indexes[i]]
        labels_count[label] = labels_count.get(label, 0) + 1

    op = operator.itemgetter(1)
    tuples = sorted(labels_count.iteritems(), key=op, reverse=True)
    return tuples[0][0]


def test():
    data_set = np.array(([0, 1], [0, 1], [1, 0], [1, 1], [2, 0], [2, 1]))
    labels = ['A', 'A', 'B', 'B', 'C', 'C']
    knn = KNN(data_set, labels, 1)
    label = knn.tell_label([1, 3])
    print label

if __name__ == '__main__':
    test()