#!/usr/bin/env python3

import numpy as np

from sklearn import svm
from sklearn.linear_model import SGDClassifier

from scatteringtree import scattering_tree
from filters.swt2 import swt2
from filters.gabortsfm import gabortsfm
from filters.gabortrafo import gabortrafo
from nonlinearity import sigmoid
from pooling import pool

from scipy.misc import imresize

from sklearn.datasets import fetch_mldata

from sklearn.decomposition import PCA, IncrementalPCA
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis

from joblib import Parallel, delayed
from multiprocessing import cpu_count

from nonlinearity import *

import cProfile

num_cpus = cpu_count()
verbose = 8
num_train = 100
num_test = 100



def mnist_to_img(data):
    img_small = data.reshape((28,28))
    return imresize(img_small, (32, 32), interp='bilinear')

# Download MNIST dataset
print('fetching MNIST dataset...')
mnist = fetch_mldata('MNIST original', data_home='./mnist_dataset')


def run_test():
    #classifier = svm.LinearSVC(multi_class='crammer_singer')
    #classifier = svm.LinearSVC()
    #classifier = SGDClassifier()
    classifier = svm.SVC(kernel='linear', cache_size=2000)

    # initialize scattering transform

    prop_filters = [
                        #[gabortsfm(5,5,32,32)],
                        #[swt2(wavelet='haar', level=3, start_level=0, frequency_decreasing_path=True)],
                        #[swt2(wavelet='haar', level=3, start_level=0, frequency_decreasing_path=True)],
                        #[swt2(wavelet='haar', level=3, start_level=0, frequency_decreasing_path=True)],
                        #[],
                        [gabortrafo(2,32,32)],
                        [gabortrafo(2,32,32)],
                        [gabortrafo(1,32,32)],
                        [],
                    ]
                    
    output_filters = [
                        [None],
                        [],
                        [],
                        [None],
                        ]
    nls = [None]*4

    pools = [
                None,
                None,
                None,
                pool((2,2), np.max),
                pool((2,2), np.max),
                pool((2,2), np.max),
                ]
                                    
    sct = scattering_tree(
                            propagation_filters=prop_filters,
                            output_filters=output_filters,
                            nonlinearities=nls,
                            poolings=pools
                            )

    # init training data

    idx = np.arange(0, int(len(mnist.data)/1))
    np.random.seed(0)
    np.random.shuffle(idx)	# We want the test data to be 'randomly' mixed.

    training_input = mnist.data[idx[0:num_train]]
    training_input = [mnist_to_img(t) for t in training_input] 	# Convert to 32x32 pixel image

    training_labels = mnist.target[idx[0:num_train]]

    # apply scattering transform to each training image
    scattered_input = Parallel(n_jobs=num_cpus, verbose=verbose)(delayed(sct.scatter_linearized)(inp) for inp in training_input)
    #scattered_input = [sct.scatter_linearized(x) for x in training_input]

    print('Length of feature vector before PCA: %d'%len(scattered_input[0]))

    # dimensionality reduction
    print('running dimensionality reduction...')
    dim_reduce = PCA(n_components=200) # reduce to n_components features
    #dim_reduce = IncrementalPCA(n_components=200, batch_size=10) # reduce to n_components features
    dim_reduce.fit(scattered_input[0:800]) 

    #dim_reduce = LinearDiscriminantAnalysis(n_components=200)
    #dim_reduce.fit(scattered_input[0:400], training_labels[0:400])

    scattered_input = dim_reduce.transform(scattered_input) # reduce training data dimensionality

    print('Length of feature vector after PCA: %d'%len(scattered_input[0]))

    # feed the classifier
    print('training SVM classifier...')
    classifier.fit(scattered_input, training_labels)

    # test

    # initialize test vectors
    test_input = mnist.data[idx[num_train:num_train+num_test]]
    test_input = [mnist_to_img(t) for t in test_input]

    # apply scattering transform to each test image
    scattered_test_input = Parallel(n_jobs=num_cpus, verbose=verbose)(delayed(sct.scatter_linearized)(inp) for inp in test_input)
    #scattered_test_input = [sct.scatter_linearized(x) for x in test_input]

    scattered_test_input = dim_reduce.transform(scattered_test_input) # reduce test data dimensionality

    test_labels = mnist.target[idx[num_train:num_train+num_test]]

    prediction = classifier.predict(scattered_test_input)

    num_correct = np.sum(prediction == test_labels)
    print("prediction accuracy: %.3f%%"%(num_correct/len(prediction)*100))
    

#cProfile.run('run_test()')
run_test()
