from torch import nn as nn

from .base_model import BaseModel
from ..nn.conv2d import DenseConv2d
from ..nn.linear import DenseLinear

__all__ = ["Conv2", "conv2", "Conv4", "conv4"]


class Conv2(BaseModel):
    def __init__(self):
        super(Conv2, self).__init__()
        self.features = nn.Sequential(DenseConv2d(1, 32, kernel_size=5, padding=2),  # 32x28x28
                                      nn.ReLU(inplace=True),
                                      nn.MaxPool2d(2, stride=2),  # 32x14x14
                                      DenseConv2d(32, 64, kernel_size=5, padding=2),  # 64x14x14
                                      nn.ReLU(inplace=True),
                                      nn.MaxPool2d(2, stride=2))  # 64x7x7

        self.classifier = nn.Sequential(DenseLinear(64 * 7 * 7, 2048),
                                        nn.ReLU(inplace=True),
                                        DenseLinear(2048, 62))
        self.collect_prunable_layers()

    def forward(self, inp):
        out = self.features(inp)
        out = out.view(out.size(0), -1)
        out = self.classifier(out)
        return out


class Conv4(BaseModel):
    def __init__(self):
        super(Conv4, self).__init__()
        self.features = nn.Sequential(DenseConv2d(3, 32, kernel_size=3, padding=1),
                                      nn.BatchNorm2d(32),
                                      nn.MaxPool2d(2),
                                      DenseConv2d(32, 32, kernel_size=3, padding=1),
                                      nn.BatchNorm2d(32),
                                      nn.MaxPool2d(2),
                                      DenseConv2d(32, 32, kernel_size=3, padding=2),
                                      nn.BatchNorm2d(32),
                                      nn.MaxPool2d(2),
                                      DenseConv2d(32, 32, kernel_size=3, padding=2),
                                      nn.BatchNorm2d(32),
                                      nn.MaxPool2d(2))

        self.classifier = DenseLinear(in_features=32 * 6 * 6, out_features=2)

    def forward(self, inp):
        out = self.features(inp)
        out = out.view(out.size(0), -1)
        out = self.classifier(out)
        return out


def conv2() -> Conv2:
    return Conv2()


def conv4() -> Conv4:
    return Conv4()

# TODO: define pretrain etc.

[((404, 425), 'torch.nn.ReLU', 'nn.ReLU', ([], {'inplace': '(True)'}), '(inplace=True)\n', (411, 425), True, 'from torch import nn as nn\n'), ((465, 490), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (['(2)'], {'stride': '(2)'}), '(2, stride=2)\n', (477, 490), True, 'from torch import nn as nn\n'), ((639, 660), 'torch.nn.ReLU', 'nn.ReLU', ([], {'inplace': '(True)'}), '(inplace=True)\n', (646, 660), True, 'from torch import nn as nn\n'), ((700, 725), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (['(2)'], {'stride': '(2)'}), '(2, stride=2)\n', (712, 725), True, 'from torch import nn as nn\n'), ((849, 870), 'torch.nn.ReLU', 'nn.ReLU', ([], {'inplace': '(True)'}), '(inplace=True)\n', (856, 870), True, 'from torch import nn as nn\n'), ((1340, 1358), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(32)'], {}), '(32)\n', (1354, 1358), True, 'from torch import nn as nn\n'), ((1398, 1413), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (['(2)'], {}), '(2)\n', (1410, 1413), True, 'from torch import nn as nn\n'), ((1538, 1556), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(32)'], {}), '(32)\n', (1552, 1556), True, 'from torch import nn as nn\n'), ((1596, 1611), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (['(2)'], {}), '(2)\n', (1608, 1611), True, 'from torch import nn as nn\n'), ((1736, 1754), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(32)'], {}), '(32)\n', (1750, 1754), True, 'from torch import nn as nn\n'), ((1794, 1809), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (['(2)'], {}), '(2)\n', (1806, 1809), True, 'from torch import nn as nn\n'), ((1934, 1952), 'torch.nn.BatchNorm2d', 'nn.BatchNorm2d', (['(32)'], {}), '(32)\n', (1948, 1952), True, 'from torch import nn as nn\n'), ((1992, 2007), 'torch.nn.MaxPool2d', 'nn.MaxPool2d', (['(2)'], {}), '(2)\n', (2004, 2007), True, 'from torch import nn as nn\n')]