#!/usr/bin/env python
"""

Naive Bayes classification on Anderson's iris data.

"""
from __future__ import division, print_function
import numpy as np
import learnml as ml
import learnml.datasets as datasets
import learnml.cross_validation as cv
try:
    import matplotlib.pyplot as plt
except ImportError:
    plt = None


if __name__ == '__main__':

    np.random.seed(0)

    # load the data
    iris = datasets.load_iris()
    print(iris)
    X, y = iris

    # split into a training and a testing set
    X_learn, y_learn, X_test, y_test = cv.split_data(X, y, frac=3/4)

    # train the (Gaussian) naive Bayes classifier
    naive = ml.NaiveBayesGaussian()
    naive.learn(X_learn, y_learn)
    print(naive)

    # make predictions on the testing set
    y_hat = naive.predict(X_test)
    print("Incorrect predictions on the testing set: {0:.2f} %"
          .format(100. * (y_hat != y_test).sum() / len(y_test)))

    if plt is not None:
        # plot learning data
        colord = {}
        for c in naive.classes:
            i = y_learn == c
            lines, = plt.plot(X_learn[i, 2], X_learn[i, 3], 'x', label=c)
            colord[c] = lines.get_color()

        # plot predictions
        for c in naive.classes:
            i = y_hat == c
            lines, = plt.plot(X_test[i, 2], X_test[i, 3], 'o', color=colord[c])

        plt.legend(loc=2)
        plt.xlabel('Petal length (cm)')
        plt.ylabel('Petal width (cm)')
        plt.savefig('iris-naive_bayes.pdf')
        plt.show()
