import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchinfo import summary

# GoogLeNet / Inception (Szegedy et al., 2014)
# 1x1 konvoluce slouží jako tzv. "bottleneck" (zúžení):
# Místo přímé 5x5 konvoluce na 192 kanálech se nejprve pomocí 1x1 konvoluce
# sníží počet kanálů ze 192 na 16. Následná prostorová 5x5 konvoluce tak počítá
# nad pouhými 16 kanály, což ušetří ~90 % parametrů i výpočetních operací.

class Block5Example(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=5, padding=2)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, X):
        out = self.relu(self.conv(X))
        return out

class Block15Example(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv_1x1_16 = nn.Conv2d(in_channels, 16, kernel_size=1)
        self.conv_5x5_32 = nn.Conv2d(16, out_channels, kernel_size=5, padding=2)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, X):
        out = self.relu(self.conv_1x1_16(X))
        out = self.relu(self.conv_5x5_32(out))
        return out


print("------------ Přímá 5x5 konvoluce (Block 5x5) -------------")
block_5 = Block5Example(192, 32)
summary(block_5, input_size=(1, 192, 28, 28), col_names=("input_size", "output_size", "num_params"))

print("------------ Redukce přes 1x1 + 5x5 konvoluci ------------")
block_15 = Block15Example(192, 32)
summary(block_15, input_size=(1, 192, 28, 28), col_names=("input_size", "output_size", "num_params"))