from numpy.linalg import norm

class KNN(object):
    def __init__(self, n_neighbours=3, classes = None):
        self.n_neighbours = n_neighbours
        self.classes = classes

    def fit(self, X_train, y_train):
        self.X_train = X_train
        self.y_train = y_train
        self.min_y = min(y_train)
        self.max_y = max(y_train)
        if not self.classes:
            self.classes = range(self.min_y, self.max_y + 1)


    def predict(self, X_test):
        predictions = []

        for x_test in X_test:
            distances = []
            for x_, y_ in zip(self.X_train, self.y_train):
                distances.append((norm(x_test - x_), y_))
            distances.sort(key=lambda x: x[0])

            cls_dict = {}
            for k in self.classes:
                cls_dict[k] = 0

            for dist, y in distances[:self.n_neighbours]:
                cls_dict[y] += 1
            predictions.append(max(cls_dict, key=cls_dict.get))

        return np.array(predictions)