#!/usr/bin/env python

# coding: utf-8

'''Subject-adaptative 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.path import join as pjoin

import h5py

import torch

import torch.nn.functional as F

from braindecode.models.deep4 import Deep4Net

from braindecode.torch_ext.optimizers import AdamW

from braindecode.torch_ext.util import set_random_seeds

from torch import nn

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

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

parser = argparse.ArgumentParser(

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

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

parser.add_argument('modelpath', type=str,

                    help='Path to the base model folder')

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

parser.add_argument('-scheme', type=int, help='Adaptation scheme', default=4)

parser.add_argument(

    '-trfrate', type=int, help='The percentage of data for adaptation', default=100)

parser.add_argument('-lr', type=float, help='Learning rate', default=0.0005)

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

args = parser.parse_args()

datapath = args.datapath

outpath = args.outpath

modelpath = args.modelpath

scheme = args.scheme

rate = args.trfrate

lr = args.lr

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

torch.cuda.set_device(args.gpu)

set_random_seeds(seed=20200205, cuda=True)

BATCH_SIZE = 16

TRAIN_EPOCH = 200

# 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]

# 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[:]

X, Y = get_data(subjs[0])

n_classes = 2

in_chans = 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=X.shape[2],

                 final_conv_length='auto').cuda()

# Deprecated.

def reset_conv_pool_block(network, block_nr):

    suffix = "_{:d}".format(block_nr)

    conv = getattr(network, 'conv' + suffix)

    kernel_size = conv.kernel_size

    n_filters_before = conv.in_channels

    n_filters = conv.out_channels

    setattr(network, 'conv' + suffix,

            nn.Conv2d(

                n_filters_before,

                n_filters,

                kernel_size,

                stride=(1, 1),

                bias=False,

            ))

    setattr(network, 'bnorm' + suffix,

            nn.BatchNorm2d(

                n_filters,

                momentum=0.1,

                affine=True,

                eps=1e-5,

            ))

    # Initialize the layers.

    conv = getattr(network, 'conv' + suffix)

    bnorm = getattr(network, 'bnorm' + suffix)

    nn.init.xavier_uniform_(conv.weight, gain=1)

    nn.init.constant_(bnorm.weight, 1)

    nn.init.constant_(bnorm.bias, 0)

def reset_model(checkpoint):

    # Load the state dict of the model.

    model.network.load_state_dict(checkpoint['model_state_dict'])

    # # Resets the last conv block

    # reset_conv_pool_block(model.network, block_nr=4)

    # reset_conv_pool_block(model.network, block_nr=3)

    # reset_conv_pool_block(model.network, block_nr=2)

    # # Resets the fully-connected layer.

    # # Parameters of newly constructed modules have requires_grad=True by default.

    # n_final_conv_length = model.network.conv_classifier.kernel_size[0]

    # n_prev_filter = model.network.conv_classifier.in_channels

    # n_classes = model.network.conv_classifier.out_channels

    # model.network.conv_classifier = nn.Conv2d(

    #     n_prev_filter, n_classes, (n_final_conv_length, 1), bias=True)

    # nn.init.xavier_uniform_(model.network.conv_classifier.weight, gain=1)

    # nn.init.constant_(model.network.conv_classifier.bias, 0)

    if scheme != 5:

        # Freeze all layers.

        for param in model.network.parameters():

            param.requires_grad = False

        if scheme in {1, 2, 3, 4}:

            # Unfreeze the FC layer.

            for param in model.network.conv_classifier.parameters():

                param.requires_grad = True

        if scheme in {2, 3, 4}:

            # Unfreeze the conv4 layer.

            for param in model.network.conv_4.parameters():

                param.requires_grad = True

            for param in model.network.bnorm_4.parameters():

                param.requires_grad = True

        if scheme in {3, 4}:

            # Unfreeze the conv3 layer.

            for param in model.network.conv_3.parameters():

                param.requires_grad = True

            for param in model.network.bnorm_3.parameters():

                param.requires_grad = True

        if scheme == 4:

            # Unfreeze the conv2 layer.

            for param in model.network.conv_2.parameters():

                param.requires_grad = True

            for param in model.network.bnorm_2.parameters():

                param.requires_grad = True

    # Only optimize parameters that requires gradient.

    optimizer = AdamW(filter(lambda p: p.requires_grad, model.network.parameters()),

                      lr=lr, weight_decay=0.5*0.001)

    model.compile(loss=F.nll_loss, optimizer=optimizer,

                  iterator_seed=20200205, )

cutoff = int(rate * 200 / 100)

# Use only session 1 data for training

assert(cutoff <= 200)

for fold, subj in enumerate(subjs):

    suffix = '_s' + str(subj) + '_f' + str(fold)

    checkpoint = torch.load(pjoin(modelpath, 'model_f' + str(fold) + '.pt'),

                            map_location='cuda:' + str(args.gpu))

    reset_model(checkpoint)

    X, Y = get_data(subj)

    X_train, Y_train = X[:cutoff], Y[:cutoff]

    X_val, Y_val = X[200:300], Y[200:300]

    X_test, Y_test = X[300:], Y[300:]

    model.fit(X_train, Y_train, epochs=TRAIN_EPOCH,

              batch_size=BATCH_SIZE, scheduler='cosine',

              validation_data=(X_val, Y_val), remember_best_column='valid_loss')

    model.epochs_df.to_csv(pjoin(outpath, 'epochs' + suffix + '.csv'))

    test_loss = model.evaluate(X_test, Y_test)

    with open(pjoin(outpath, 'test' + suffix + '.json'), 'w') as f:

        json.dump(test_loss, f)

dfile.close()