#!/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.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 nonlinearity import *

from visualize import visualize_tree



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')


if __name__ == '__main__':
    # initialize scattering transform

    prop_filters = [
                        #[swt2(wavelet='db1', level=2, start_level=1, frequency_decreasing_path=True)],
                        #[swt2(wavelet='db1', level=2, start_level=1, frequency_decreasing_path=True)],
                        #[swt2(wavelet='db1', level=2, start_level=0, frequency_decreasing_path=True)],
                        #[gabortrafo(2,32,32)],
                        [gabortrafo(2,32,32)],
                        [gabortrafo(1,32,32)],
                        #[swt2(wavelet='haar', level=2, start_level=0, frequency_decreasing_path=True)],
                        [],
                    ]
                    
    output_filters = [
                        [],
                        [None],
                        [],
                        [],
                        ]
    nls = [
        np.abs,
        np.abs,
        None,
        None,
    ]

    pools = [
                #pool((2,2), np.max),
                None,
                None,
                None,
                ]
                                    
    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]

    signal = mnist_to_img(training_input[1])
    visualize_tree(sct, signal)

    
