"""
Mean log loss from 5-fold CV: 0.480065955962
"""
import copy
import itertools
import numpy as np
import lasagne
import math
import os
import theano
import theano.tensor as T
import time

from lasagne.layers import DenseLayer, DropoutLayer, InputLayer, get_all_params
from lasagne.nonlinearities import rectify, softmax
from lasagne.objectives import categorical_crossentropy, Objective
from lasagne.updates import adagrad

from sklearn import feature_extraction
from sklearn.base import BaseEstimator
from sklearn.cross_validation import StratifiedShuffleSplit
from sklearn.utils import check_random_state

from otto_utils import consts, utils


MODEL_NAME = 'model_09_nn_adagrad'
MODE = 'cv'  # cv|submission|holdout|tune


class NeuralNetwork(BaseEstimator):
    def __init__(self, n_hidden=20, max_epochs=150, batch_size=200,
                 lr=0.01, epsilon=0.9, dropout=0.5, valid_ratio=0.0,
                 use_valid=False, verbose=0, random_state=None):
        self.n_hidden = n_hidden
        self.max_epochs = max_epochs
        self.batch_size = batch_size
        self.lr = lr
        self.epsilon = epsilon
        self.dropout = dropout
        self.valid_ratio = valid_ratio
        self.use_valid = use_valid
        self.verbose = verbose
        self.random_state = random_state
        # State
        self.score_ = None
        self.classes_ = None
        self.n_classes_ = None
        self.model = None

    def fit(self, data, targets, sample_weight=None):
        self.classes_, indices = np.unique(targets, return_inverse=True)
        self.n_classes_ = self.classes_.shape[0]

        random_state = check_random_state(self.random_state)

        # Shuffle data and eventually split on train and validation sets
        if self.valid_ratio > 0:
            strat_shuffled_split = StratifiedShuffleSplit(targets, test_size=self.valid_ratio,
                                                          n_iter=1, random_state=self.random_state)
            train_index, valid_index = [s for s in strat_shuffled_split][0]
            X_train, y_train = data[train_index], targets[train_index]
            X_valid, y_valid = data[valid_index], targets[valid_index]
        else:
            X_train, y_train = data, targets
            X_valid, y_valid = np.array([]), np.array([])

        if self.verbose > 5:
            print 'X_train: %s, y_train: %s' % (X_train.shape, y_train.shape)
            if self.use_valid:
                print 'X_valid: %s, y_valid: %s' % (X_valid.shape, y_valid.shape)

        # Prepare theano variables
        dataset = dict(
            X_train=theano.shared(lasagne.utils.floatX(X_train)),
            y_train=T.cast(theano.shared(y_train), 'int32'),
            X_valid=theano.shared(lasagne.utils.floatX(X_valid)),
            y_valid=T.cast(theano.shared(y_valid), 'int32'),
            num_examples_train=X_train.shape[0],
            num_examples_valid=X_valid.shape[0],
            input_dim=X_train.shape[1],
            output_dim=self.n_classes_,
        )

        if self.verbose > 0:
            print "Building model and compiling functions..."
        output_layer = self.build_model(dataset['input_dim'])
        iter_funcs = self.create_iter_functions(dataset, output_layer)

        if self.verbose > 0:
            print "Starting training..."
        now = time.time()
        results = []
        try:
            for epoch in self.train(iter_funcs, dataset, output_layer):
                if self.verbose > 1:
                    print "Epoch {} of {} took {:.3f}s".format(
                        epoch['number'], self.max_epochs, time.time() - now)
                now = time.time()
                results.append([epoch['number'], epoch['train_loss'], epoch['valid_loss']])
                if self.verbose > 1:
                    print "  training loss:\t\t{:.6f}".format(epoch['train_loss'])
                    print "  validation loss:\t\t{:.6f}".format(epoch['valid_loss'])
                    print "  validation accuracy:\t\t{:.2f} %%".format(
                        epoch['valid_accuracy'] * 100)

                if epoch['number'] >= self.max_epochs:
                    break

            if self.verbose > 0:
                print 'Minimum validation error: %f (epoch %d)' % \
                      (epoch['best_val_error'], epoch['best_val_iter'])

        except KeyboardInterrupt:
            pass

        return self

    def predict(self, data):
        preds, _ = self.make_predictions(data)

        return preds

    def predict_proba(self, data):
        _, proba = self.make_predictions(data)

        return proba

    def score(self):
        return self.score_

    # Private methods
    def build_model(self, input_dim):
        l_in = InputLayer(shape=(self.batch_size, input_dim))

        l_hidden1 = DenseLayer(l_in, num_units=self.n_hidden, nonlinearity=rectify)
        l_hidden1_dropout = DropoutLayer(l_hidden1, p=self.dropout)

        l_hidden2 = DenseLayer(l_hidden1_dropout, num_units=self.n_hidden, nonlinearity=rectify)
        l_hidden2_dropout = DropoutLayer(l_hidden2, p=self.dropout)

        l_hidden3 = DenseLayer(l_hidden2_dropout, num_units=self.n_hidden, nonlinearity=rectify)
        l_hidden3_dropout = DropoutLayer(l_hidden3, p=self.dropout)