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

try:
    from IPython import display
    display.set_matplotlib_formats('svg')
except (ImportError, AttributeError):
    pass

# https://arxiv.org/abs/1409.1556v6
# VGG-11 (Konfigurace A z článku)
# Celkem 11 váhových vrstev: 8 konvolučních + 3 plně propojené (FC)

class MyVGG11(nn.Module):
    def __init__(self, in_channels=3, num_classes=1000):
        super().__init__()
        
        # Blok 1: 1x Conv (64) + MaxPool
        self.conv1 = nn.Conv2d(in_channels, 64, kernel_size=3, padding=1)

        # Blok 2: 1x Conv (128) + MaxPool

        # Blok 3: 2x Conv (256, 256) + MaxPool

        # Blok 4: 2x Conv (512, 512) + MaxPool

        # Blok 5: 2x Conv (512, 512) + MaxPool


        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # Klasifikátor: 3x FC (4096 -> 4096 -> num_classes)
        self.flatten = nn.Flatten()
        #self.fc1 = nn.Linear(512 * 7 * 7, 4096)
        # Alternativně: self.fc1 = nn.LazyLinear(4096)
        self.fc1 = nn.LazyLinear(4096)
        self.fc2 = nn.Linear(4096, 4096)
        self.fc3 = nn.Linear(4096, num_classes)
        self.dropout = nn.Dropout(p=0.5)

    def forward(self, x):
        print("input:", x.shape)

        # Blok 1
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        print("block1 (pool):", x.shape)

        # Blok 2


        # Blok 3


        # Blok 4


        # Blok 5


        # Klasifikátor
        x = self.flatten(x)
        print("flatten:", x.shape)

        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        print("fc1:", x.shape)

        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        print("fc2:", x.shape)

        x = self.fc3(x)
        print("fc3 (výstup):", x.shape)

        return x

device = "cpu"
model = MyVGG11().to(device)

test_tensor = torch.randn(1, 3, 224, 224).to(device)
model(test_tensor)

print("Struktura modelu:")
print(model)

print("Torchinfo přehled pro MyVGG11:")
summary(model, input_size=(1, 3, 224, 224), col_names=("input_size", "output_size", "num_params"), device=device)

# Srovnání s referenční implementací v torchvision
torch_vgg11 = torchvision.models.vgg11(weights=None).to(device)
print("Torchinfo přehled pro torchvision VGG-11:")
summary(torch_vgg11, input_size=(1, 3, 224, 224), col_names=("input_size", "output_size", "num_params"), device=device)

# Varianta s Batch Normalization
torch_vgg11_bn = torchvision.models.vgg11_bn(weights=None).to(device)
print("Torchinfo přehled pro torchvision VGG-11 s BatchNorm:")
summary(torch_vgg11_bn, input_size=(1, 3, 224, 224), col_names=("input_size", "output_size", "num_params"), device=device)