from sklearn.neighbors import NearestNeighbors, KNeighborsClassifier
import numpy as np


def main():
    # a test of this method using an arbitrarily generated list of 5 vectors with 3 features each
    # nearestNeighbors([[1, 1, 0], [1, 0, 0], [0, 0, 0], [0, 5, 5]], [[1, 1, 4]])
    kNearestNeighbors([[1, 1, 0], [1, 0, 0], [0, 0, 0], [0, 5, 5]], ["three", 2, 3, "5"], [[1, 1, 0], [0, 5, 5]])


def kNearestNeighbors(data: list, classifications: list, test_data: list):
    kn = KNeighborsClassifier(n_neighbors=2)
    kn.fit(data, classifications)
    print(kn.predict(test_data))


def nearestNeighbors(data: list, test_data: list):
    x = np.array(data)
    nbrs = NearestNeighbors(n_neighbors=1, algorithm='ball_tree').fit(x)
    dist, indicies = nbrs.kneighbors(test_data)
    print("Indicies:")
    print(indicies)
    print("Distances:")
    print(dist)
    return indicies, dist


if __name__ == '__main__':
    main()