import numpy as np
import torch as th
from torch import nn
from torch.nn import init
from torch.nn.functional import elu

from braindecode.models.base import BaseModel
from braindecode.torch_ext.init import glorot_weight_zero_bias
from braindecode.torch_ext.modules import Expression
from braindecode.torch_ext.util import np_to_var

class Conv2dWithConstraint(nn.Conv2d):
    def __init__(self, *args, max_norm=1, **kwargs):
        self.max_norm = max_norm
        super(Conv2dWithConstraint, self).__init__(*args, **kwargs)

    def forward(self, x):
        self.weight.data = th.renorm(self.weight.data, p=2, dim=0,
                                     maxnorm=self.max_norm)
        return super(Conv2dWithConstraint, self).forward(x)

class EEGNetv4(BaseModel):
    """
    EEGNet v4 model from [EEGNet4]_.

    Notes
    -----
    This implementation is not guaranteed to be correct, has not been checked
    by original authors, only reimplemented from the paper description.

    References
    ----------

    .. [EEGNet4] Lawhern, V. J., Solon, A. J., Waytowich, N. R., Gordon,
       S. M., Hung, C. P., & Lance, B. J. (2018).
       EEGNet: A Compact Convolutional Network for EEG-based
       Brain-Computer Interfaces.
       arXiv preprint arXiv:1611.08024.
    """

    def __init__(self, in_chans,
                 n_classes,
                 final_conv_length='auto',
                 input_time_length=None,
                 pool_mode='mean',
                 F1=8,
                 D=2,
                 F2=16, # usually set to F1*D (?)
                 kernel_length=64,
                 third_kernel_size=(8, 4),
                 drop_prob=0.25
                 ):

        if final_conv_length == 'auto':
            assert input_time_length is not None
        self.__dict__.update(locals())
        del self.self

    def create_network(self):
        pool_class = dict(max=nn.MaxPool2d, mean=nn.AvgPool2d)[self.pool_mode]
        model = nn.Sequential()
        # b c 0 1
        # now to b 1 0 c
        model.add_module('dimshuffle', Expression(_transpose_to_b_1_c_0))

        model.add_module('conv_temporal', nn.Conv2d(
            1, self.F1, (1, self.kernel_length), stride=1, bias=False,
            padding=(0, self.kernel_length // 2,)))
        model.add_module('bnorm_temporal', nn.BatchNorm2d(
            self.F1, momentum=0.01, affine=True, eps=1e-3), )
        model.add_module('conv_spatial', Conv2dWithConstraint(
            self.F1, self.F1 * self.D, (self.in_chans, 1), max_norm=1, stride=1, bias=False,
            groups=self.F1,
            padding=(0, 0)))

        model.add_module('bnorm_1', nn.BatchNorm2d(
            self.F1 * self.D, momentum=0.01, affine=True, eps=1e-3), )
        model.add_module('elu_1', Expression(elu))

        model.add_module('pool_1', pool_class(
            kernel_size=(1, 4), stride=(1, 4)))
        model.add_module('drop_1', nn.Dropout(p=self.drop_prob))

        # https://discuss.pytorch.org/t/how-to-modify-a-conv2d-to-depthwise-separable-convolution/15843/7
        model.add_module('conv_separable_depth', nn.Conv2d(
            self.F1 * self.D, self.F1 * self.D, (1, 16), stride=1, bias=False, groups=self.F1 * self.D,
            padding=(0, 16 // 2)))
        model.add_module('conv_separable_point', nn.Conv2d(
            self.F1 * self.D, self.F2, (1, 1), stride=1, bias=False,
            padding=(0, 0)))

        model.add_module('bnorm_2', nn.BatchNorm2d(
            self.F2, momentum=0.01, affine=True, eps=1e-3), )
        model.add_module('elu_2', Expression(elu))
        model.add_module('pool_2', pool_class(
            kernel_size=(1, 8), stride=(1, 8)))
        model.add_module('drop_2', nn.Dropout(p=self.drop_prob))

        out = model(np_to_var(np.ones(
            (1, self.in_chans, self.input_time_length, 1),
            dtype=np.float32)))
        n_out_virtual_chans = out.cpu().data.numpy().shape[2]

        if self.final_conv_length == 'auto':
            n_out_time = out.cpu().data.numpy().shape[3]
            self.final_conv_length = n_out_time

        model.add_module('conv_classifier', nn.Conv2d(
            self.F2, self.n_classes,
            (n_out_virtual_chans, self.final_conv_length,), bias=True))
        model.add_module('softmax', nn.LogSoftmax())
        # Transpose back to the the logic of braindecode,
        # so time in third dimension (axis=2)
        model.add_module('permute_back', Expression(_transpose_1_0))
        model.add_module('squeeze', Expression(_squeeze_final_output))

        glorot_weight_zero_bias(model)
        return model


def _transpose_to_b_1_c_0(x):
    return x.permute(0, 3, 1, 2)


def _transpose_1_0(x):
    return x.permute(0, 1, 3, 2)


# remove empty dim at end and potentially remove empty time dim
# do not just use squeeze as we never want to remove first dim
def _squeeze_final_output(x):
    assert x.size()[3] == 1
    x = x[:, :, :, 0]
    if x.size()[2] == 1:
        x = x[:, :, 0]
    return x



class EEGNet(object):
    """
    EEGNet model from [EEGNet]_.

    Notes
    -----
    This implementation is not guaranteed to be correct, has not been checked
    by original authors, only reimplemented from the paper description.

    References
    ----------

    .. [EEGNet] Lawhern, V. J., Solon, A. J., Waytowich, N. R., Gordon, 
       S. M., Hung, C. P., & Lance, B. J. (2016).
       EEGNet: A Compact Convolutional Network for EEG-based
       Brain-Computer Interfaces.
       arXiv preprint arXiv:1611.08024.
    """
    def __init__(self, in_chans,
                 n_classes,
                 final_conv_length='auto',
                 input_time_length=None,
                 pool_mode='max',
                 second_kernel_size=(2,32),
                 third_kernel_size=(8,4),
                 drop_prob=0.25
                 ):

        if final_conv_length == 'auto':
            assert input_time_length is not None
        self.__dict__.update(locals())
        del self.self

    def create_network(self):
        pool_class = dict(max=nn.MaxPool2d, mean=nn.AvgPool2d)[self.pool_mode]
        model = nn.Sequential()
        n_filters_1 = 16
        model.add_module('conv_1', nn.Conv2d(
            self.in_chans, n_filters_1, (1, 1), stride=1, bias=True))
        model.add_module('bnorm_1', nn.BatchNorm2d(
            n_filters_1, momentum=0.01, affine=True, eps=1e-3),)
        model.add_module('elu_1', Expression(elu))
        # transpose to examples x 1 x (virtual, not EEG) channels x time
        model.add_module('permute_1', Expression(lambda x: x.permute(0,3,1,2)))

        model.add_module('drop_1', nn.Dropout(p=self.drop_prob))

        n_filters_2 = 4
        # keras padds unequal padding more in front, so padding
        # too large should be ok.
        # Not padding in time so that croped training makes sense
        # https://stackoverflow.com/questions/43994604/padding-with-even-kernel-size-in-a-convolutional-layer-in-keras-theano

        model.add_module('conv_2', nn.Conv2d(
            1, n_filters_2, self.second_kernel_size, stride=1,
            padding=(self.second_kernel_size[0] // 2, 0),
            bias=True))
        model.add_module('bnorm_2',nn.BatchNorm2d(
            n_filters_2, momentum=0.01, affine=True, eps=1e-3),)
        model.add_module('elu_2', Expression(elu))
        model.add_module('pool_2', pool_class(
            kernel_size=(2, 4), stride=(2, 4)))
        model.add_module('drop_2', nn.Dropout(p=self.drop_prob))

        n_filters_3 = 4
        model.add_module('conv_3', nn.Conv2d(
            n_filters_2, n_filters_3, self.third_kernel_size, stride=1,
            padding=(self.third_kernel_size[0] // 2, 0),
            bias=True))
        model.add_module('bnorm_3',nn.BatchNorm2d(
            n_filters_3, momentum=0.01, affine=True, eps=1e-3),)
        model.add_module('elu_3', Expression(elu))
        model.add_module('pool_3', pool_class(
            kernel_size=(2, 4), stride=(2, 4)))
        model.add_module('drop_3', nn.Dropout(p=self.drop_prob))



        out = model(np_to_var(np.ones(
            (1, self.in_chans, self.input_time_length, 1),
            dtype=np.float32)))
        n_out_virtual_chans = out.cpu().data.numpy().shape[2]

        if self.final_conv_length == 'auto':
            n_out_time = out.cpu().data.numpy().shape[3]
            self.final_conv_length = n_out_time

        model.add_module('conv_classifier', nn.Conv2d(
            n_filters_3, self.n_classes,
            (n_out_virtual_chans, self.final_conv_length,), bias=True))
        model.add_module('softmax', nn.LogSoftmax())
        # Transpose back to the the logic of braindecode,
        # so time in third dimension (axis=2)
        model.add_module('permute_2', Expression(lambda x: x.permute(0,1,3,2)))
        model.add_module('squeeze',  Expression(_squeeze_final_output))
        glorot_weight_zero_bias(model)
        return model
