# -*- coding: utf-8 -*-
"""
Created on Thu Jul 30 21:28:32 2015.

Script written by Tim Hochberg with parameter tweaks by Bluefool.
https://www.kaggle.com/bitsofbits/grasp-and-lift-eeg-detection/naive-nnet

Modifications: rc, alex
"""
import os
import sys
if __name__ == '__main__' and __package__ is None:
    filePath = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    sys.path.append(filePath)

import yaml
from glob import glob
import numpy as np
import pandas as pd
from time import time

from sklearn.metrics import roc_auc_score

# Lasagne (& friends) imports
import theano
from nolearn.lasagne import BatchIterator, NeuralNet, TrainSplit
from lasagne.objectives import aggregate, binary_crossentropy
from lasagne.layers import (InputLayer, DropoutLayer, DenseLayer, Conv1DLayer,
                            Conv2DLayer)
from lasagne.updates import nesterov_momentum, adam
from theano.tensor.nnet import sigmoid

from mne import concatenate_raws, pick_types

from preprocessing.aux import creat_mne_raw_object
from preprocessing.filterBank import FilterBank

# Silence some warnings from lasagne
import warnings
warnings.filterwarnings('ignore', '.*topo.*')
warnings.filterwarnings('ignore', module='.*lasagne.init.*')
warnings.filterwarnings('ignore', module='.*nolearn.lasagne.*')

####
yml = yaml.load(open(sys.argv[1]))
fileName = yml['Meta']['file']

filt2Dsize = yml['filt2Dsize'] if 'filt2Dsize' in yml.keys() else 0
filters = yml['filters']
delay = yml['delay']
skip = yml['skip']

if 'bags' in yml.keys():
    bags = yml['bags']
else:
    bags = 3

mode = sys.argv[2]
if mode == 'val':
    test = False
elif mode == 'test':
    test = True
else:
    raise('Invalid mode. Please specify either val or test')


###########
SUBJECTS = list(range(1, 13))
TRAIN_SERIES = list(range(1, 9))
TEST_SERIES = [9, 10]

N_ELECTRODES = 32
N_EVENTS = 6

SAMPLE_SIZE = delay
DOWNSAMPLE = skip
TIME_POINTS = SAMPLE_SIZE // DOWNSAMPLE

TRAIN_SIZE = 5120

# We encapsulate the event / electrode data in a Source object.


def preprocessData(data):
    """Preprocess data with filterbank."""
    fb = FilterBank(filters)
    return fb.transform(data)


class Source:

    """Loads, preprocesses and holds data."""

    mean = None
    std = None

    def load_raw_data(self, subject, series):
        """Load data for a subject / series."""
        test = series == TEST_SERIES
        if not test:
            fnames = [glob('../data/train/subj%d_series%d_data.csv' %
                      (subject, i)) for i in series]
        else:
            fnames = [glob('../data/test/subj%d_series%d_data.csv' %
                      (subject, i)) for i in series]
        fnames = list(np.concatenate(fnames))
        fnames.sort()
        raw_train = [creat_mne_raw_object(fname, read_events=not test)
                     for fname in fnames]
        raw_train = concatenate_raws(raw_train)
        # pick eeg signal
        picks = pick_types(raw_train.info, eeg=True)

        self.data = raw_train._data[picks].transpose()

        self.data = preprocessData(self.data)

        if not test:

            self.events = raw_train._data[32:].transpose()

    def normalize(self):
        """normalize data."""
        self.data -= self.mean
        self.data /= self.std


class TrainSource(Source):

    """Source for training data."""

    def __init__(self, subject, series_list):
        """Init."""
        self.load_raw_data(subject, series_list)
        self.mean = self.data.mean(axis=0)
        self.std = self.data.std(axis=0)
        self.normalize()


# Note that Test/Submit sources use the mean/std from the training data.
# This is both standard practice and avoids using future data in theano
# test set.

class TestSource(Source):

    """Source for test data."""

    def __init__(self, subject, series, train_source):
        """Init."""
        self.load_raw_data(subject, series)
        self.mean = train_source.mean
        self.std = train_source.std
        self.normalize()


# Lay out the Neural net.


class LayerFactory:

    """Helper class that makes laying out Lasagne layers more pleasant."""

    def __init__(self):
        """Init."""
        self.layer_cnt = 0
        self.kwargs = {}

    def __call__(self, layer, layer_name=None, **kwargs):
        """Call."""
        self.layer_cnt += 1
        name = layer_name or "layer{0}".format(self.layer_cnt)
        for k, v in kwargs.items():
            self.kwargs["{0}_{1}".format(name, k)] = v
        return (name, layer)


class IndexBatchIterator(BatchIterator):

    """Generate BatchData from indices.

    Rather than passing the data into the fit function, instead we just pass in
    indices to the data.  The actual data is then grabbed from a Source object
    that is passed in at the creation of the IndexBatchIterator. Passing in a
    '-1' grabs a random value from the Source.

    As a result, an "epoch" here isn't a traditional epoch, which looks at all
    the time points. Instead a random subsamle of 0.8*TRAIN_SIZE points from
    the training data are used each "epoch" and 0.2 TRAIN_SIZE points are use
    for validation.

    """

    def __init__(self, source, *args, **kwargs):
        """Init."""
        super(IndexBatchIterator, self).__init__(*args, **kwargs)
        self.source = source
        if source is not None:
            # Tack on (SAMPLE_SIZE-1) copies of the first value so that it is
            # easy to grab
            # SAMPLE_SIZE POINTS even from the first location.
            x = source.data
            input_shape = [len(x) + (SAMPLE_SIZE - 1), N_ELECTRODES]
            self.augmented = np.zeros(input_shape, dtype=np.float32)
            self.augmented[SAMPLE_SIZE-1:] = x
            self.augmented[:SAMPLE_SIZE-1] = x[0]
        if filt2Dsize:
            input_shape = [self.batch_size, 1, N_ELECTRODES, TIME_POINTS]
            self.Xbuf = np.zeros(input_shape, np.float32)
        else:
            input_shape = [self.batch_size, N_ELECTRODES, TIME_POINTS]
            self.Xbuf = np.zeros(input_shape, np.float32)
        self.Ybuf = np.zeros([self.batch_size, N_EVENTS], np.float32)

    def transform(self, X_indices, y_indices):
        """Transform."""
        X_indices, y_indices = super(IndexBatchIterator,
                                     self).transform(X_indices, y_indices)
        [count] = X_indices.shape
        # Use preallocated space
        X = self.Xbuf[:count]
        Y = self.Ybuf[:count]
        for i, ndx in enumerate(X_indices):
            if ndx == -1:
                ndx = np.random.randint(len(self.source.events))
            sample = self.augmented[ndx:ndx+SAMPLE_SIZE]
            # Reverse so we get most recent point, otherwise downsampling drops
            # the last
            # DOWNSAMPLE-1 points.
            if filt2Dsize:
                X[i][0] = sample[::-1][::DOWNSAMPLE].transpose()
            else:
                X[i] = sample[::-1][::DOWNSAMPLE].transpose()

            if y_indices is not None:
                Y[i] = self.source.events[ndx]
        Y = None if (y_indices is None) else Y
        return X, Y


# Simple / Naive net. Borrows from Daniel Nouri's Facial Keypoint Detection
# Tutorial

def create_net(train_source, test_source, batch_size=128, max_epochs=100,
               train_val_split=False):
    """Create NN."""
    if train_val_split:
        train_val_split = TrainSplit(eval_size=0.2)
    else:
        train_val_split = TrainSplit(eval_size=False)

    batch_iter_train = IndexBatchIterator(train_source, batch_size=batch_size)
    batch_iter_test = IndexBatchIterator(test_source, batch_size=batch_size)
    LF = LayerFactory()

    dense = 1024  # larger (1024 perhaps) would be better

    if filt2Dsize:
        inputLayer = LF(InputLayer, shape=(None, 1, N_ELECTRODES, TIME_POINTS))
        convLayer = LF(Conv2DLayer, num_filters=8, filter_size=(N_ELECTRODES, filt2Dsize))
    else:
        inputLayer = LF(InputLayer, shape=(None, N_ELECTRODES, TIME_POINTS))
        convLayer = LF(Conv1DLayer, num_filters=8, filter_size=1)

    layers = [
        inputLayer,
        LF(DropoutLayer, p=0.5),
        convLayer,
        # Standard fully connected net from now on
        LF(DenseLayer, num_units=dense),
        LF(DropoutLayer, p=0.5),
        LF(DenseLayer, num_units=dense),
        LF(DropoutLayer, p=0.5),
        LF(DenseLayer, layer_name="output", num_units=N_EVENTS,
           nonlinearity=sigmoid)
    ]

    def loss(x, t):
        return aggregate(binary_crossentropy(x, t))

    if filt2Dsize:
        nnet = NeuralNet(y_tensor_type=theano.tensor.matrix,
                         layers=layers,
                         batch_iterator_train=batch_iter_train,
                         batch_iterator_test=batch_iter_test,
                         max_epochs=max_epochs,
                         verbose=0,
                         update=adam,
                         update_learning_rate=0.001,
                         objective_loss_function=loss,
                         regression=True,
                         train_split=train_val_split,
                         **LF.kwargs)
    else:
        nnet = NeuralNet(y_tensor_type=theano.tensor.matrix,
                         layers=layers,
                         batch_iterator_train=batch_iter_train,
                         batch_iterator_test=batch_iter_test,
                         max_epochs=max_epochs,
                         verbose=0,
                         update=nesterov_momentum,
                         update_learning_rate=0.02,
                         update_momentum=0.9,
                         # update=adam,
                         # update_learning_rate=0.001,
                         objective_loss_function=loss,
                         regression=True,
                         train_split=train_val_split,
                         **LF.kwargs)

    return nnet


# Do the training.
print 'Running in mode %s, saving to file %s' % (mode,fileName)
report = pd.DataFrame(index=[fileName])
start_time = time()

train_indices = np.zeros([TRAIN_SIZE], dtype=int) - 1

np.random.seed(67534)

valid_series = [7, 8]
max_epochs = 100

if test is False:
    probs_bags = []
    for bag in range(bags):
        probs_tot = []
        lbls_tot = []
        for subject in range(1, 13):
            tseries = sorted(set(TRAIN_SERIES) - set(valid_series))
            train_source = TrainSource(subject, tseries)
            test_source = TestSource(subject, valid_series, train_source)
            net = create_net(train_source, test_source, max_epochs=max_epochs,
                             train_val_split=False)
            dummy = net.fit(train_indices, train_indices)
            indices = np.arange(len(test_source.data))
            probs = net.predict_proba(indices)
    
            auc = np.mean([roc_auc_score(trueVals, p) for trueVals, p in
                          zip(test_source.events.T, probs.T)])
            print 'Bag %d, subject %d, AUC: %.5f' % (bag, subject, auc)
            probs_tot.append(probs)
            lbls_tot.append(test_source.events)
    
        probs_tot = np.concatenate(probs_tot)
        lbls_tot = np.concatenate(lbls_tot)
        auc = np.mean([roc_auc_score(trueVals, p) for trueVals, p in
                      zip(lbls_tot.transpose(), probs_tot.transpose())])
        print auc
        probs_bags.append(probs_tot)
    
    probs_bags = np.mean(probs_bags, axis=0)
    np.save('val/val_%s.npy' % fileName, [probs_bags])

else:
    probs_bags = []
    for bag in range(bags):
        probs_tot = []
        for subject in range(1, 13):
            tseries = set(TRAIN_SERIES)
            train_source = TrainSource(subject, tseries)
            test_source = TestSource(subject, TEST_SERIES, train_source)
            net = create_net(train_source, test_source, max_epochs=max_epochs,
                             train_val_split=False)
            dummy = net.fit(train_indices, train_indices)
            indices = np.arange(len(test_source.data))
            probs = net.predict_proba(indices)
            print 'Bag %d, subject %d' % (bag, subject)
            probs_tot.append(probs)
        probs_tot = np.concatenate(probs_tot)
        probs_bags.append(probs_tot)
    
    probs_bags = np.mean(probs_bags, axis=0)
    np.save('test/test_%s.npy' % fileName, [probs_bags])


prefix = 'test_' if test else 'val_'
end_time = time()
report['Time'] = end_time - start_time
report.to_csv("report/%s_%s.csv" % (prefix, fileName))
print report