import torch
import torch.nn as nn

from .utils import Conv2dWithNorm, LinearWithNorm
from ..tools.docs import fill_doc


__all__ = ["DeepConvNet"]


@fill_doc
class DeepConvNet(nn.Module):
    """Deep Learning With Convolutional Neural Networks for EEG Decoding and
    Visualization (Deep ConvNet).

    Deep ConvNet [1]_ had four convolution-max-pooling blocks, with a special
    first block designed to handle EEG input, followed by three standard
    convolutionmax-pooling blocks and a dense softmax classification layer.
    The first convolutional block was split into two layers in order to better
    handle the large number of input channelsone input channel per electrode
    compared to three input channels (one per color) in rgb-images. In the
    first layer, each filter performs a convolution over time, and in the
    second layer, each filter performs a spatial filtering with weights for all
    possible pairs of electrodes with filters of the preceding temporal
    convolution. Note that as there is no activation function in between the
    two layers, they could in principle be combined into one layer. Using two
    layers however implicitly regularizes the overall convolution by forcing a
    separation of the linear transformation into a combination of a temporal
    convolution and a spatial filter.

    Parameters
    ----------
    %(nCh)s
    %(nTime)s
    %(nCls)s
    dropout : float
        Dropout rate.

    References
    ----------
    .. [1] R. T. Schirrmeister et al., “Deep learning with convolutional neural
        networks for EEG decoding and visualization,” Human Brain Mapping,
        vol. 38, no. 11, pp.5391-5420, 2017, doi: 10.1002/hbm.23730.
    """

    def __init__(self, nCh: int, nTime: int, nCls: int, dropout: float = 0.25) -> None:
        super().__init__()
        self.nCh = nCh
        self.nTime = nTime
        kernel_size = [1, 10]
        filter_layer = [25, 50, 100, 200]

        first_layer = nn.Sequential(
            Conv2dWithNorm(1, 25, kernel_size, max_norm=2),
            Conv2dWithNorm(25, 25, (nCh, 1), bias=False, max_norm=2),
            nn.BatchNorm2d(25),
            nn.ELU(),
            nn.MaxPool2d((1, 3)),
        )
        middle_layer = nn.Sequential(
            *[
                nn.Sequential(
                    nn.Dropout(dropout),
                    Conv2dWithNorm(in_f, out_f, kernel_size),
                    nn.BatchNorm2d(out_f),
                    nn.ELU(),
                    nn.MaxPool2d((1, 3)),
                )
                for in_f, out_f in zip(filter_layer, filter_layer[1:])
            ]
        )
        self.conv_layer = nn.Sequential(first_layer, middle_layer)

        linear_in = self._forward_flatten().shape[1]
        self.head = nn.Sequential(
            nn.Flatten(),
            LinearWithNorm(linear_in, nCls, max_norm=0.5),
            nn.LogSoftmax(dim=1),
        )

    def _forward_flatten(self):
        x = torch.rand(1, 1, self.nCh, self.nTime)
        x = self.conv_layer(x)
        x = torch.flatten(x, start_dim=1, end_dim=-1)
        return x

    def forward(self, x):
        """Forward pass function that processes the input EEG data and produces
        the decoded results.

        Parameters
        ----------
        x : Tensor
            Input EEG data, shape `(batch_size, 1, nCh, nTime)`.

        Returns
        -------
        cls_prob : Tensor
            Predicted class probability, shape `(batch_size, cls)`.
        """
        x = self.conv_layer(x)
        x = self.head(x)
        return x
