import pandas as pd
import numpy as np

import theano
import theano.tensor as T

import lasagne

import time
import sys

# based on Lasagne's tutorial

# log-softmax and log-domain cross-entropy functions
# might be required if using ReLUs
def log_softmax(x):
    xdev = x - x.max(1, keepdims=True)
    return xdev - T.log(T.sum(T.exp(xdev), axis=1, keepdims=True))

def categorical_crossentropy_logdomain(log_predictions, targets):
    return -T.sum(targets * log_predictions, axis=1)

def load_dataset():
    # The competition datafiles are in the directory ../input
    # Read competition data files:
    train = pd.read_csv("../input/train.csv")
    test  = pd.read_csv("../input/test.csv")
    
    pixel_labels = ['pixel' + str(i) for i in range(784)]
    X_train = train[pixel_labels].values[:]
    y_train = train['label'].values[:]
    
    X_train = X_train.reshape((-1, 1, 28, 28))
    y_train = y_train.astype(np.uint8)
    
    X_test = test[pixel_labels].values[:]
    X_test = X_test.reshape((-1, 1, 28, 28))
    
    return X_train, y_train, X_test
    
def build_mlp(input_var=None):
    """
    2 x 800 -> softmax - 10
    20% drop at 0, 50% drop on hidden
    activation function might have to be bounded (i.e. not ReLU),
    because it seems that it breaks gradients (resulting in NaNs)
    """
    l_in = lasagne.layers.InputLayer(shape=(None, 1, 28, 28), input_var=input_var)
    l_in_drop = lasagne.layers.DropoutLayer(l_in, p=0.2)
    l_hid1 = lasagne.layers.DenseLayer(
        l_in_drop, num_units=800,
        nonlinearity=lasagne.nonlinearities.sigmoid,
        W=lasagne.init.GlorotUniform())
    l_hid1_drop = lasagne.layers.DropoutLayer(l_hid1, p=0.5)
    l_hid2 = lasagne.layers.DenseLayer(
        l_hid1_drop, num_units=800,
        nonlinearity=lasagne.nonlinearities.sigmoid,
        W=lasagne.init.GlorotUniform())
    l_hid2_drop = lasagne.layers.DropoutLayer(l_hid2, p=0.5)
    l_out = lasagne.layers.DenseLayer(
        l_hid2_drop, num_units=10,
        nonlinearity=lasagne.nonlinearities.softmax)
        
    return l_out

def build_cnn(input_var=None):
    # As a third model, we'll create a CNN of two convolution + pooling stages
    # and a fully-connected hidden layer in front of the output layer.

    network = lasagne.layers.InputLayer(shape=(None, 1, 28, 28),
                                        input_var=input_var)

    # Convolutional layer with 32 kernels of size 5x5.
    network = lasagne.layers.Conv2DLayer(
            network, num_filters=32, filter_size=(5, 5),
            nonlinearity=lasagne.nonlinearities.rectify,
            W=lasagne.init.GlorotUniform())
    network = lasagne.layers.MaxPool2DLayer(network, pool_size=(2, 2))

    # Another convolution with 32 5x5 kernels, and another 2x2 pooling:
    network = lasagne.layers.Conv2DLayer(
            network, num_filters=32, filter_size=(5, 5),
            nonlinearity=lasagne.nonlinearities.rectify)
    network = lasagne.layers.MaxPool2DLayer(network, pool_size=(2, 2))

    # A fully-connected layer of 256 units with 50% dropout on its inputs:
    network = lasagne.layers.DenseLayer(
            lasagne.layers.dropout(network, p=.5),
            num_units=256,
            nonlinearity=lasagne.nonlinearities.rectify)

    # And, finally, the 10-unit output layer with 50% dropout on its inputs:
    network = lasagne.layers.DenseLayer(
            lasagne.layers.dropout(network, p=.5),
            num_units=10,
            nonlinearity=lasagne.nonlinearities.softmax)

    return network
    
def iterate_minibatches(inputs, targets, batchsize, shuffle=False):
    assert len(inputs) == len(targets)
    if shuffle:
        indices = np.arange(len(inputs))
        np.random.shuffle(indices)
    for start_idx in range(0, len(inputs) - batchsize + 1, batchsize):
        if shuffle:
            excerpt = indices[start_idx:start_idx + batchsize]
        else:
            excerpt = slice(start_idx, start_idx + batchsize)
        yield inputs[excerpt], targets[excerpt]

model = 'cnn'
print('Loading data...')
X_train, y_train, X_test = load_dataset()

input_var = T.tensor4('inputs')
target_var = T.ivector('targets')

print("Building model and compiling functions...")
if model is 'mlp':
    network = build_mlp(input_var)
elif model is 'cnn':
    network = build_cnn(input_var)
else:
    print('unknown model')
    sys.exit()
prediction = lasagne.layers.get_output(network)
loss = lasagne.objectives.categorical_crossentropy(prediction, target_var)
loss = loss.mean()

params = lasagne.layers.get_all_params(network, trainable=True)
updates = lasagne.updates.nesterov_momentum(
    loss, params, learning_rate=0.01, momentum=0.9)

test_prediction = lasagne.layers.get_output(network, deterministic=True)
#test_loss = lasagne.objectives.categorical_crossentropy(test_prediction,
#    target_var)
#test_loss = test_loss.mean()
#test_acc = T.mean(T.eq(T.argmax(test_prediction, axis=1), target_var),
#    dtype=theano.config.floatX)

train_fn = theano.function([input_var, target_var], loss, updates=updates)
pred_fn = theano.function([input_var], test_prediction)

print("Starting training...")
num_epochs=2
for epoch in range(num_epochs):
    # In each epoch, we do a full pass over the training data:
    train_err = 0
    train_batches = 0
    start_time = time.time()
    for batch in iterate_minibatches(X_train, y_train, 500, shuffle=True):
        inputs, targets = batch
        train_err += train_fn(inputs, targets)
        train_batches += 1
    print("Epoch {} of {} took {:.3f}s".format(
        epoch + 1, num_epochs, time.time() - start_time))
    print("  training loss:\t\t{:.6f}".format(train_err / train_batches))

answers = pred_fn(X_test)
answers = list(map(lambda pr: np.argmax(pr), answers))

d = { 'Label': answers }
frame = pd.DataFrame(data=d, index=[i for i in range(1, len(X_test)+1)])
with open('answers.csv', mode='w') as ans_file:
    frame.to_csv(ans_file, index_label='ImageId')