import numpy as np
import torch
import matplotlib.pyplot as plt
import torch.nn as nn
import torch.nn.functional as F
import cv2 as cv
import torchvision
from torchinfo import summary

# https://www.robots.ox.ac.uk/~vgg/research/very_deep/
# =====================================================================
# Srovnání podle článku VGG (Simonyan & Zisserman, 2014):
# Proč použít 3 konvoluční vrstvy 3x3 místo jedné vrstvy 7x7?
#
# 1. Stejné efektivní zorné pole (receptive field): 3 vrstvy 3x3 pokrývají 7x7.
# 2. Více nelinearit: 3x ReLU místo 1x ReLU (výraznější rozhodovací funkce).
# 3. Méně parametrů (pro C vstupních i výstupních kanálů, bez biasu):
#    - Blok 3x (3x3): 3 * (3^2 * C^2) = 27 * C^2
#      Pro C = 32: 27 * 32^2 = 27 * 1024 = 27 648 parametrů
#    - Blok 1x (7x7): 1 * (7^2 * C^2) = 49 * C^2
#      Pro C = 32: 49 * 32^2 = 49 * 1024 = 50 176 parametrů
#    - Jedna vrstva 7x7 má o (49 - 27) / 27 ≈ 81 % více parametrů
# =====================================================================

# 3x 3x3 vs. 1x 7x7
class Block3x3x3Example(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1, bias=False)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False)
        self.conv3 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1, bias=False)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, X):
        out = self.relu(self.conv1(X))
        out = self.relu(self.conv2(out))
        out = self.relu(self.conv3(out))
        return out

print("------------Block 3x 3x3 Example-----------------")
block_333 = Block3x3x3Example(32, 32)
summary(block_333, input_size=(1, 32, 28, 28), col_names=("input_size", "output_size", "num_params"))

class Block1x7x7Example(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=7, padding=3, bias=False)
        self.relu = nn.ReLU(inplace=True)

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

print("------------Block 1x 7x7 Example-----------------")
block_177 = Block1x7x7Example(32, 32)
summary(block_177, input_size=(1, 32, 28, 28), col_names=("input_size", "output_size", "num_params"))