#!/bin/env python3

import numpy as np

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 skimage.io import imshow, show

import matplotlib.pyplot as plt

from sklearn.datasets import fetch_mldata

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


for data in mnist.data[40000:40001]:
    img_small = data.reshape((28,28))
    img = np.zeros((32,32))
    img[0:28, 0:28] = img_small
    # normalize image
    img = img / np.max(img)

    mask_width = 0.8
        
    prop_filters = [
                    #[simple_swt2(directions=['h', 'v', 'd', 'a'], mask_width=mask_width, num_scales=1)],
                    #[simple_swt2(directions=['h', 'v', 'd', 'a'], mask_width=mask_width, num_scales=1)],
                    #[swt2(wavelet='haar', level=3, start_level=0)],
                    #[gabortsfm(5,5,32,32)],
                    #[gabortsfm(3,3,32,32)],
                    #[gabortsfm(3,3,16,16)],
                    #[swt2(wavelet='haar', level=3, start_level=0)],
                    [gabortrafo(2,32,32)],
                    [gabortrafo(1,32,32)],
                    [],
                    ]
    output_filters = [
                        [],
                        [None],
                        []
                        ]
    nls = [None]*3
    pools = [None, pool((2,2), np.max), pool((1,1), np.max)]
                                    
    sct = scattering_tree(
                            propagation_filters=prop_filters,
                            output_filters=output_filters,
                            nonlinearities=nls,
                            poolings=pools
                            )
                            
                            
    scattered = sct.scatter(img)
    
#	print(sct.scatter_linearized(img))

    # plot
    fig, axes = plt.subplots(5, 5, figsize=(10, 10), subplot_kw={'xticks': [], 'yticks': []})

    fig.subplots_adjust(hspace=0.3, wspace=0.05)


    for ax, s in zip(axes.flat, scattered):
        im = ax.imshow(s, interpolation='nearest')

        #plt.colorbar(im)
    plt.show()
