"""
This is a script that is used to generate a confussion matrix for
a classification method. This uses 10-k cross_validation with in
order to provide sensible resutls and not overfit.
"""

__author__ = "Andres Ruiz"
__license__ = "Apache"
__email__ = "afruizc __thingy__ cs unm edu"

import numpy as np
from sklearn.cross_validation import KFold
from sklearn.metrics import confusion_matrix, accuracy_score, log_loss
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

import svm_bow

def plot_confusion_matrix(cm, title='Confusion matrix', normalized=True,
                            cmap=plt.cm.Oranges, save_file=""):
    """
    Displays the confussion matrix indicated by `cm`. If argument
    `normalized` is Ture, then the matrix is normalized. Optionally
    the image can be saved to a file

    Arguments:
    ----------
    `cm`: The confusion matrix to be displayed.
    `title`: The title for the window.
    `normalized`: If True, normalizes the matrix before showing it.
    `cmap`: Colormap to use.
    `save_file`: If string different than empty, the resulting image is
    stored in such file.
    """
    if normalized:
        cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
    plt.imshow(cm, interpolation='nearest', cmap=cmap)
    plt.title(title)
    plt.colorbar()
    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')

    if save_file:
        plt.savefig(save_file)

def get_indices(data, indices):
    result = []
    for i in indices:
        result.append(data[i])
    return result

def main():
    e = svm_bow.Executor()
    e.load_data()
    e.config_model()

    fold = KFold(len(e.train['data']), n_folds=10)
    conf_mat_avg = np.zeros((9, 9))
    c = 0
    for train, test in fold:
        X_train = get_indices(e.train['data'], train)
        X_test = get_indices(e.train['data'], test)
        y_train = get_indices(e.train['target'], train)
        y_test = get_indices(e.train['target'], test)
        c += 1
        print("Fitting run {}.".format(c))
        model = e.param_tunning.fit(X_train, y_train)
        print("Predicting...")
        y_pred = model.predict(X_test)
        y_pred_prob = model.predict_proba(X_test)
        conf_matrix = confusion_matrix(y_test, y_pred)
        accruacy = accuracy_score(y_test, y_pred)
        loss = log_loss(y_test, y_pred_prob)
        plot_confusion_matrix(conf_matrix,
                              save_file='fold_{}.png'.format(c))
        np.savetxt('conf_matrix_fold{}'.format(c), conf_matrix)
        print("Fold %d. Accuracy: %lf Loss: %lf" % (c, accruacy, loss))
        conf_mat_avg += conf_matrix

    np.savetxt('conf_matrix.txt', conf_mat_avg)

    conf_mat_avg /= 10.0
    plot_confusion_matrix(conf_mat_avg, save_file='final_cm.png')

if __name__ == '__main__':
    main()
