import numpy as np
%matplotlib inline
import matplotlib
import matplotlib.pyplot as plt
from sklearn.cross_validation import train_test_split
from sklearn.datasets import load_digits
from sklearn.preprocessing import OneHotEncoder, StandardScaler
#
from feather import ann, stacker

# Get the digits data dump
digits = load_digits()

# Extract X and y arrays
X = digits.data
y = digits.target

# Scale X
sc = StandardScaler()
sc.fit(X)
Xs = sc.transform(X)

# Onehot-encode the target labels
enc = OneHotEncoder()
enc.fit([[l] for l in set(y)])
onehot_labels = enc.transform(y.reshape(-1, 1)).toarray()

# Split into train and test samples
train_X, test_X, train_y, test_y =  \
train_test_split(Xs, onehot_labels, test_size=0.2)

# Initialize network. The input dimension is 64
# (each image is an 8x8 array), and here we assign
# hidden depths of [20,15]. Thus, the full stack of
# layers depths is [64, 20, 15, 10], where the 10-dimensional
# output represents a probability distribution over the
# ten possible digit labels.
mlp = ann.ANN(
    train_X,
    train_y=train_y,
    hidden_depths=[20,15],
    test_X=test_X,
    test_y=test_y,
    eta=0.03,
    lamb=5.0,
    batch_size=30,
    initzero=False,
    activation_type='tanh',
    softmax=True,
    cost_type='Xentropy')

# Fit over some epochs
mlp.fit(40)
# If you notice that the train accuracy is much better
# than the test, then you could adjust for overfitting by
# increasing the lambda parameter and fit over more epochs.
# You can also decrease the learning parameter to fine-tune
# fitting:
# mlp.lamb = 10.
# mlp.eta = 0.01
# mlp.fit(30)

# Get predictions for any test set. These are
# probabilities for the labels [0,1,2,3,4,5,6,7,8,9]
softmax_preds = mlp.predict_proba(test_X)
print softmax_preds[0]

# Intialize autoencoder, where the deepest enconding
# dimension is 2. Since the input dimension is 64,
# the full stack of layer depths becomes [64, 20, 9, 2, 9, 20, 64].
# Essentially, we are trying to reconstruct the 64-dimensional
# input data by passing it through a "narrow" encoding funnel
# of dimension 2.
ae = stacker.Autoencoder(Xs,
    hidden_depths=[20,9,2],
    epochs=50,
    eta=0.03,
    lamb=5,
    batch_size=30,
    activation_type='tanh')

# Train autoencoder
ae.run_stacker()
# Fine-train autoencoder, now that it has been stacked together
ae.stacked_net.fit(30)

# Get encoded representation of the test set
test_X_encoded = stacker.encode(test_X, ae.stacked_net)

# Get the labels for the test set, these are the digit numbers
test_labels = np.argmax(test_y, axis=1)

# Plot
plt.figure(figsize=(14,9))
for i in range(test_X_encoded.shape[0]):
    plt.text(test_X_encoded[i, 0], test_X_encoded[i, 1], str(test_labels[i]),
             color=plt.cm.Set1(test_labels[i] / 10.),
             fontdict={'weight': 'bold', 'size': 9})
plt.xticks([])
plt.yticks([])
plt.xlim([-1.1, 1.1])
plt.ylim([-1.1, 1.1])
plt.show()