#!/usr/bin/env python

# coding: utf-8

'''Subject-independent classification with KU Data,

using Deep ConvNet model from [1].

References

----------

.. [1] Schirrmeister, R. T., Springenberg, J. T., Fiederer, L. D. J.,

   Glasstetter, M., Eggensperger, K., Tangermann, M., Hutter, F. & Ball, T. (2017).

   Deep learning with convolutional neural networks for EEG decoding and

   visualization.

   Human Brain Mapping , Aug. 2017. Online: http://dx.doi.org/10.1002/hbm.23730

'''

import argparse

import json

import logging

import sys

from os import makedirs

from os.path import join as pjoin

from shutil import copy2, move

import h5py

import numpy as np

import torch

import torch.nn.functional as F

from braindecode.datautil.signal_target import SignalAndTarget

from braindecode.models.deep4 import Deep4Net

from braindecode.torch_ext.optimizers import AdamW

from braindecode.torch_ext.util import set_random_seeds

from sklearn.model_selection import KFold

logging.basicConfig(format='%(asctime)s %(levelname)s : %(message)s',

                    level=logging.INFO, stream=sys.stdout)

parser = argparse.ArgumentParser(

    description='Subject-independent classification with KU Data')

parser.add_argument('datapath', type=str, help='Path to the h5 data file')

parser.add_argument('outpath', type=str, help='Path to the result folder')

parser.add_argument('-fold', type=int,

                    help='k-fold index, starts with 0', required=True)

parser.add_argument('-gpu', type=int, help='The gpu device to use', default=0)

args = parser.parse_args()

datapath = args.datapath

outpath = args.outpath

fold = args.fold

assert(fold >= 0 and fold < 54)

# Randomly shuffled subject.

subjs = [35, 47, 46, 37, 13, 27, 12, 32, 53, 54, 4, 40, 19, 41, 18, 42, 34, 7,

         49, 9, 5, 48, 29, 15, 21, 17, 31, 45, 1, 38, 51, 8, 11, 16, 28, 44, 24,

         52, 3, 26, 39, 50, 6, 23, 2, 14, 25, 20, 10, 33, 22, 43, 36, 30]

test_subj = subjs[fold]

cv_set = np.array(subjs[fold+1:] + subjs[:fold])

kf = KFold(n_splits=6)

dfile = h5py.File(datapath, 'r')

torch.cuda.set_device(args.gpu)

set_random_seeds(seed=20200205, cuda=True)

BATCH_SIZE = 16

TRAIN_EPOCH = 200  # consider 200 for early stopping

# Get data from single subject.

def get_data(subj):

    dpath = '/s' + str(subj)

    X = dfile[pjoin(dpath, 'X')]

    Y = dfile[pjoin(dpath, 'Y')]

    return X, Y

def get_multi_data(subjs):

    Xs = []

    Ys = []

    for s in subjs:

        x, y = get_data(s)

        Xs.append(x[:])

        Ys.append(y[:])

    X = np.concatenate(Xs, axis=0)

    Y = np.concatenate(Ys, axis=0)

    return X, Y

cv_loss = []

for cv_index, (train_index, test_index) in enumerate(kf.split(cv_set)):

    train_subjs = cv_set[train_index]

    valid_subjs = cv_set[test_index]

    X_train, Y_train = get_multi_data(train_subjs)

    X_val, Y_val = get_multi_data(valid_subjs)

    X_test, Y_test = get_data(test_subj)

    train_set = SignalAndTarget(X_train, y=Y_train)

    valid_set = SignalAndTarget(X_val, y=Y_val)

    test_set = SignalAndTarget(X_test[200:], y=Y_test[200:])

    n_classes = 2

    in_chans = train_set.X.shape[1]

    # final_conv_length = auto ensures we only get a single output in the time dimension

    model = Deep4Net(in_chans=in_chans, n_classes=n_classes,

                     input_time_length=train_set.X.shape[2],

                     final_conv_length='auto').cuda()

    # these are good values for the deep model

    optimizer = AdamW(model.parameters(), lr=1*0.01, weight_decay=0.5*0.001)

    model.compile(loss=F.nll_loss, optimizer=optimizer, iterator_seed=1, )

    # Fit the base model for transfer learning, use early stopping as a hack to remember the model

    exp = model.fit(train_set.X, train_set.y, epochs=TRAIN_EPOCH, batch_size=BATCH_SIZE, scheduler='cosine',

                    validation_data=(valid_set.X, valid_set.y), remember_best_column='valid_loss')

    rememberer = exp.rememberer

    base_model_param = {

        'epoch': rememberer.best_epoch,

        'model_state_dict': rememberer.model_state_dict,

        'optimizer_state_dict': rememberer.optimizer_state_dict,

        'loss': rememberer.lowest_val

    }

    torch.save(base_model_param, pjoin(

        outpath, 'model_f{}_cv{}.pt'.format(fold, cv_index)))

    model.epochs_df.to_csv(

        pjoin(outpath, 'epochs_f{}_cv{}.csv'.format(fold, cv_index)))

    cv_loss.append(rememberer.lowest_val)

    test_loss = model.evaluate(test_set.X, test_set.y)

    with open(pjoin(outpath, 'test_base_s{}_f{}_cv{}.json'.format(test_subj, fold, cv_index)), 'w') as f:

        json.dump(test_loss, f)

best_cv = np.argmin(cv_loss)

best_dir = pjoin(outpath, "best")

makedirs(best_dir, exist_ok=True)

with open(pjoin(best_dir, "fold_bestcv.txt"), 'a') as f:

    f.write("{}, {}\n".format(fold, best_cv))

copy2(pjoin(outpath, 'model_f{}_cv{}.pt'.format(fold, best_cv)),

      pjoin(best_dir, 'model_f{}.pt'.format(fold)))

copy2(pjoin(outpath, 'epochs_f{}_cv{}.csv'.format(fold, best_cv)),

      pjoin(best_dir, 'epochs_f{}.csv'.format(fold)))

copy2(pjoin(outpath, 'test_base_s{}_f{}_cv{}.json'.format(test_subj, fold, best_cv)),

      pjoin(best_dir, 'test_base_s{}_f{}.json'.format(test_subj, fold)))

dfile.close()