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', (['inputs', '(32)', '(5)'], {'stride': '(1)', 'padding': '(2)'}), 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', ([], {'kernel_size': '(3)', 'stride': '(2)'}), True, 'import torch.nn as nn\n'), (16, 'utils.BBBlayers.BBBConv2d', 'BBBConv2d', (['(32)', '(64)', '(5)'], {'stride': '(1)', 'padding': '(2)'}), 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', ([], {'kernel_size': '(3)', 'stride': '(2)'}), True, 'import torch.nn as nn\n'), (20, 'utils.BBBlayers.BBBConv2d', 'BBBConv2d', (['(64)', '(128)', '(5)'], {'stride': '(1)', 'padding': '(1)'}), 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', ([], {'kernel_size': '(3)', 'stride': '(2)'}), True, 'import torch.nn as nn\n'), (24, 'utils.BBBlayers.FlattenLayer', 'FlattenLayer', (['(2 * 2 * 128)'], {}), False, 'from utils.BBBlayers import BBBConv2d, BBBLinearFactorial, FlattenLayer\n'), (25, 'utils.BBBlayers.BBBLinearFactorial', 'BBBLinearFactorial', (['(2 * 2 * 128)', '(1000)'], {}), 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', (['(1000)', '(1000)'], {}), 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', (['(1000)', 'outputs'], {}), False, 'from utils.BBBlayers import BBBConv2d, BBBLinearFactorial, FlattenLayer\n'), (37, 'torch.nn.ModuleList', 'nn.ModuleList', (['layers'], {}), True, 'import torch.nn as nn\n')]