"""TensorViz teaching example: a small, untrained image classifier."""
import torch
import torch.nn as nn


class ImageFeatures(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 8, 3, padding=1)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2)
        self.conv2 = nn.Conv2d(8, 16, 3, padding=1)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)

    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.pool1(x)
        x = self.conv2(x)
        x = self.relu2(x)
        return self.pool2(x)


# tensorviz: input=1,3,32,32
class SmallCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = ImageFeatures()
        self.flatten = nn.Flatten()
        self.head = nn.Linear(1024, 10)

    def forward(self, x):
        x = self.features(x)
        x = self.flatten(x)
        return self.head(x)


if __name__ == "__main__":
    torch.manual_seed(0)
    with torch.no_grad():
        output = SmallCNN().eval()(torch.ones(1, 3, 32, 32))
    assert tuple(output.shape) == (1, 10)
    print("Input [1, 3, 32, 32] -> output", list(output.shape))