import numpy as np
from sklearn.neighbors import KNeighborsClassifier
from sklearn.cross_validation import train_test_split, cross_val_score
from sklearn import svm
from sklearn.decomposition import PCA


def load_data(filename):
    return np.genfromtxt(filename, delimiter=',')


def perform_pca(Xs, n_components):
    # FIXME: TOTALLY WRONG! use fit() to train and then
    # transform
    print("Performing PCA with %s components" % (str(n_components),))
    pca = PCA(n_components=n_components)
    return [pca.fit_transform(X) for X in Xs]


def find_best_svm(X, y):
    print("Searching SVMs")
    scores = []
    for kernel in ['linear', 'rbf', 'sigmoid']:
        print(kernel)
        # For some reason, poly needed really long or
        # did even get stuck, so removed it
        clf = svm.SVC(kernel=kernel, C=1.0)
        score = cross_val_score(clf, X, y, cv=5)
        print("%s: %0.2f" % (kernel, score.mean()))

        scores.append((score.mean(), kernel))

    return max(scores)


def semi_supervised_set(clf, X, y, X_test, thresh=0.9):
    """
    Create a semi-supervised training set
    using estimations from classifier clf which will be
    trained on X and y before, but only take
    guesses where the probability is at least thresh
    """
    clf.fit(X, y)
    y_test = clf.predict(X_test)

    new_X = np.zeros((X.shape[0] + X_test.shape[0], X.shape[1]))
    new_y = np.zeros(X.shape[0] + X_test.shape[0])
    new_X[0:X.shape[0], :] = X
    new_y[0:X.shape[0]] = y

    added_count = 0
    for i, probs in enumerate(clf.predict_proba(X_test)):
        if probs[y_test[i]] > thresh:
            new_X[X.shape[0] + added_count, :] = X_test[i, :]
            new_y[X.shape[0] + added_count] = y_test[i]

            added_count += 1

    new_X.resize((X.shape[0] + added_count, X.shape[1]))
    new_y.resize(X.shape[0] + added_count)

    return new_X, new_y

def use_classifier(clf, X, y, X_test, filename):
    """
    Perform classification onto test data and save the
    result into filename
    """
    clf.fit(X, y)

    res = clf.predict(X_test)
    output = np.zeros((res.shape[0], 2), dtype='int')
    output[:, 0] = np.arange(1, res.shape[0] + 1)
    output[:, 1] = res
    np.savetxt(filename, output, fmt='%d', delimiter=',',
               header='Id,Solution', comments='')


X = load_data('train.csv')
y = load_data('trainLabels.csv')
X_test = load_data('test.csv')

scores_svm = []
for components in range(1, 2):  # removed mle, due to bad result
    #X_pca, X_test_pca = perform_pca((X, X_test), components)

    #score, kernel = find_best_svm(X_pca, y)
    score, kernel = find_best_svm(X, y)
    scores_svm.append((round(score, 2), components, kernel))

score, components, kernel = max(scores_svm, key=lambda x: (x[0], -x[1]))
print("Using %s kernel with %s components due to score %0.2f" %
      (kernel, str(components), score))
#X_pca, X_test_pca = perform_pca((X, X_test), components)
clf = svm.SVC(kernel=kernel, C=1.0, probability=True)
for i in range(2):
    X, y = semi_supervised_set(clf, X, y, X_test, 0.9)
clf = svm.SVC(kernel=kernel, C=1.0, probability=True)
#use_classifier(clf, X_pca, y, X_test_pca, 'submission_svm.csv')
use_classifier(clf, X, y, X_test, 'submission_svm_semi.csv')
