import torch.nn as nn
import torch

class ConvBlock(nn.Module):

    def __init__(self,
                 in_channels,
                 out_channels,
                 kernel_size,
                 max_pool=True,
                 max_pool_factor=1.0,
                 batch_norm=False,
                 dropout=False,
                 mi=None,
                 activation='relu'):
        super(ConvBlock, self).__init__()
        stride = (int(2 * max_pool_factor), int(2 * max_pool_factor))
        if max_pool:
            self.max_pool = nn.MaxPool2d(kernel_size=stride,
                                         stride=stride,
                                         ceil_mode=False,
                                         )
            stride = (1, 1)
        else:
            self.max_pool = lambda x: x

        if batch_norm:
            self.normalize = nn.BatchNorm2d(out_channels,affine=True)
        else:
            self.normalize = nn.Identity()

        if dropout:
            self.drop = nn.Dropout(p=0.5)
        else:
            self.drop = nn.Identity()

        if activation == 'relu':
            self.actv = nn.ReLU()
        else:
            self.actv = nn.Tanh()

        self.conv = nn.Conv2d(in_channels,
                              out_channels,kernel_size,stride=stride,padding=1)

        self.mi = mi

    def forward(self, x, i, get_mi, noise):
        mi = None
        x = self.conv(x)
        x = self.normalize(x)
        x = self.actv(x)
        x = x + torch.rand_like(x)*noise

        if get_mi and self.mi:
            media, covarianza = self.mi.get_stats(x.clone())
            mi = self.mi.multilevelMI(media, covarianza)
            
        x = self.drop(x)
        x = self.max_pool(x)

        return x, mi

class ConvBase(nn.Sequential):
    def __init__(self,
                 filters,
                 channels=3,
                 max_pool=False,
                 max_pool_factor=1.0,
                 batch_norm=True,
                 dropout=False,
                 mi=None,
                 activation='relu'):
        core = [ConvBlock(channels,
                          filters[0],
                          (3, 3),
                          max_pool=max_pool,
                          max_pool_factor=max_pool_factor,
                          batch_norm=batch_norm,
                          dropout=dropout,
                          mi=mi,
                          activation=activation),
                ]

        for l in range(1,len(filters)):
            core.append(ConvBlock(filters[l-1],
                                  filters[l],
                                  kernel_size=(3, 3),
                                  max_pool=max_pool,
                                  max_pool_factor=max_pool_factor,
                                  batch_norm=batch_norm,
                                  dropout=dropout,
                                  mi=mi,
                                  activation=activation))
        super(ConvBase, self).__init__(*core)

class Net(nn.Module):
    def __init__(self, num_filters, output_size, channels=3, mi=None, batch_norm=True, dropout=False, noise=1e-4, activation='relu'):
        super(Net, self).__init__()
        if len(num_filters) == 4:
            factor = 4
        else:
            factor = 1

        self.mi = mi
        self.base = ConvBase(filters=num_filters,
                             channels=channels,
                             max_pool=False,
                             max_pool_factor=4 // 4,
                             batch_norm=batch_norm,
                             dropout=dropout,
                             mi=self.mi,
                             activation=activation)
        self.linear = nn.Linear(factor * num_filters[-1], output_size, bias=True)
        
        self.noise = noise

    def forward(self, x, get_mi=False):
        mi = {}

        for i,elem in enumerate(self.base):
            x, mi_layer = elem(x, i, get_mi, self.noise)
            if mi_layer:
                mi[i] = mi_layer
            #     media, covarianza = self.mi.get_stats(x.clone())
            #     mi[i] = self.mi.multilevelMI(media, covarianza)

        x = self.linear(x.view(x.size(0), -1))

        x = x + torch.rand_like(x)*self.noise
        if self.mi and get_mi:
            media, covarianza = self.mi.get_stats_lin(x.clone())
            mi[i+1] = self.mi.multilevelMI(media, covarianza)
            #mi[i+1] = self.mi.mutual_information(media, covarianza)

        return x, mi