import matplotlib
import matplotlib.pyplot as pyplot
import numpy
from matplotlib.colors import ListedColormap
from sklearn import neighbors, datasets


def get_iris_dataset():
    iris = datasets.load_iris()
    return iris.data[:, :2], iris.target


def get_knn_classifier(X, y, n_neighbors=None):
    if not n_neighbors:
        n_neighbors = 6
    classifier = neighbors.KNeighborsClassifier(n_neighbors, weights='distance')
    classifier.fit(X, y)
    return classifier, n_neighbors


def get_meshgrid(X, y, h=None):
    if not h:
        h = .02
    x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
    y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
    return numpy.meshgrid(
        numpy.arange(x_min, x_max, h), numpy.arange(y_min, y_max, h))


def predict(classifier, mesh_xx, mesh_yy):
    Z = classifier.predict(numpy.c_[mesh_xx.ravel(), mesh_yy.ravel()])
    return Z.reshape(mesh_xx.shape)


def plot_classified_regions(X, y, classifier, n_neighbors):
    xx, yy = get_meshgrid(X, y)
    Z = predict(classifier, xx, yy)
    pyplot.figure()
    pyplot.pcolormesh(xx, yy, Z)

    # Plot also the training points
    cmap = ListedColormap(['#FFAAAA', '#AAFFAA','#00AAFF'])
    pyplot.scatter(X[:, 0], X[:, 1], c=y, cmap=cmap, alpha=0.8)

    pyplot.xlim(xx.min(), xx.max())
    pyplot.ylim(yy.min(), yy.max())
    pyplot.title("3-Class classification (k = %i)" % (n_neighbors))
    pyplot.savefig('knearest.png')


X, y = get_iris_dataset()
knn, n_neighbors = get_knn_classifier(X, y)
plot_classified_regions(X, y, knn, n_neighbors)
