import torch.nn as nn
from utils.BBBlayers import BBBConv2d, BBBLinearFactorial, FlattenLayer

class BBB3Conv3FC(nn.Module):
    """

    Simple Neural Network having 3 Convolution
    and 3 FC layers with Bayesian layers.
    """
    def __init__(self, outputs, inputs):
        super(BBB3Conv3FC, self).__init__()
        self.conv1 = BBBConv2d(inputs, 32, 5, stride=1, padding=2)
        self.soft1 = nn.Softplus()
        self.pool1 = nn.MaxPool2d(kernel_size=3, stride=2)

        self.conv2 = BBBConv2d(32, 64, 5, stride=1, padding=2)
        self.soft2 = nn.Softplus()
        self.pool2 = nn.MaxPool2d(kernel_size=3, stride=2)

        self.conv3 = BBBConv2d(64, 128, 5, stride=1, padding=1)
        self.soft3 = nn.Softplus()
        self.pool3 = nn.MaxPool2d(kernel_size=3, stride=2)

        self.flatten = FlattenLayer(2 * 2 * 128)
        self.fc1 = BBBLinearFactorial(2 * 2 * 128, 1000)
        self.soft5 = nn.Softplus()

        self.fc2 = BBBLinearFactorial(1000, 1000)
        self.soft6 = nn.Softplus()

        self.fc3 = BBBLinearFactorial(1000, outputs)

        layers = [self.conv1, self.soft1, self.pool1, self.conv2, self.soft2, self.pool2,
                  self.conv3, self.soft3, self.pool3, self.flatten, self.fc1, self.soft5,
                  self.fc2, self.soft6, self.fc3]

        self.layers = nn.ModuleList(layers)

    def probforward(self, x):
        'Forward pass with Bayesian weights'
        kl = 0
        for layer in self.layers:
            if hasattr(layer, 'convprobforward') and callable(layer.convprobforward):
                x, _kl, = layer.convprobforward(x)
                kl += _kl

            elif hasattr(layer, 'fcprobforward') and callable(layer.fcprobforward):
                x, _kl, = layer.fcprobforward(x)
                kl += _kl
            else:
                x = layer(x)
        logits = x
        return logits, kl
[(12, 'utils.BBBlayers.BBBConv2d', 'BBBConv2d', False, 'from utils.BBBlayers import BBBConv2d, BBBLinearFactorial, FlattenLayer\n'), (13, 'torch.nn.Softplus', 'nn.Softplus', True, 'import torch.nn as nn\n'), (14, 'torch.nn.MaxPool2d', 'nn.MaxPool2d', True, 'import torch.nn as nn\n'), (16, 'utils.BBBlayers.BBBConv2d', 'BBBConv2d', False, 'from utils.BBBlayers import BBBConv2d, BBBLinearFactorial, FlattenLayer\n'), (17, 'torch.nn.Softplus', 'nn.Softplus', True, 'import torch.nn as nn\n'), (18, 'torch.nn.MaxPool2d', 'nn.MaxPool2d', True, 'import torch.nn as nn\n'), (20, 'utils.BBBlayers.BBBConv2d', 'BBBConv2d', False, 'from utils.BBBlayers import BBBConv2d, BBBLinearFactorial, FlattenLayer\n'), (21, 'torch.nn.Softplus', 'nn.Softplus', True, 'import torch.nn as nn\n'), (22, 'torch.nn.MaxPool2d', 'nn.MaxPool2d', True, 'import torch.nn as nn\n'), (24, 'utils.BBBlayers.FlattenLayer', 'FlattenLayer', False, 'from utils.BBBlayers import BBBConv2d, BBBLinearFactorial, FlattenLayer\n'), (25, 'utils.BBBlayers.BBBLinearFactorial', 'BBBLinearFactorial', False, 'from utils.BBBlayers import BBBConv2d, BBBLinearFactorial, FlattenLayer\n'), (26, 'torch.nn.Softplus', 'nn.Softplus', True, 'import torch.nn as nn\n'), (28, 'utils.BBBlayers.BBBLinearFactorial', 'BBBLinearFactorial', False, 'from utils.BBBlayers import BBBConv2d, BBBLinearFactorial, FlattenLayer\n'), (29, 'torch.nn.Softplus', 'nn.Softplus', True, 'import torch.nn as nn\n'), (31, 'utils.BBBlayers.BBBLinearFactorial', 'BBBLinearFactorial', False, 'from utils.BBBlayers import BBBConv2d, BBBLinearFactorial, FlattenLayer\n'), (37, 'torch.nn.ModuleList', 'nn.ModuleList', True, 'import torch.nn as nn\n')]