#!/bin/env python3

import numpy as np

from scatteringtree import scatter_single_node
from filters.swt2 import swt2
from nonlinearity import sigmoid
from pooling import pool

from skimage import data
from skimage.io import imshow, show
from skimage.transform import resize

import matplotlib.pyplot as plt

from sklearn.datasets import fetch_mldata

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

def mnist_to_img(data):
    img_small = data.reshape((28,28))
    return resize(img_small, (32, 32))

for data in mnist.data[0:1]:
    
    img = mnist_to_img(data)

    # normalize image
    img = img / np.max(img)

    prop_filters = [
                        swt2(wavelet='haar', level=3, start_level=0)
                    ]
    output_filters = []

    pooling = None

    scattered = scatter_single_node(img, 
                                    propagation_filters=prop_filters,
                                    output_filters=output_filters,
                                    nonlinearity=sigmoid,
                                    pooling=pooling)

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

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


    for ax, s in zip(axes.flat, scattered['prop']):
        im = ax.imshow(s)
        
    plt.show()
