import numpy as np
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F

from torchinfo import summary


# https://papers.nips.cc/paper_files/paper/2012/hash/c399862d3b9d6b76c8436e924a68c45b-Abstract.html

class MyAlexNet(nn.Module):
    def __init__(self, num_classes=1000):
        super(MyAlexNet, self).__init__()
        self.conv1 = nn.Conv2d(3, 96, kernel_size=11, stride=4, padding=2)
        self.conv2 = nn.Conv2d(96, 256, kernel_size=5, stride=1, padding=2)
        self.conv3 = nn.Conv2d(256, 384, kernel_size=3, stride=1, padding=1)
        self.conv4 = nn.Conv2d(384, 384, kernel_size=3, stride=1, padding=1)
        self.conv5 = nn.Conv2d(384, 256, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size = 3, stride = 2)
        # self.linear1 = nn.Linear(9216, 4096)
        # Pokud se mění vstupní rozlišení nebo pooling/strides, je vhodné použít nn.LazyLinear pro automatické odvození velikosti.
        self.flatten = nn.Flatten()
        self.linear1 = nn.LazyLinear(4096)
        self.linear2 = nn.Linear(4096, 4096)
        self.linear3 = nn.Linear(4096, num_classes)
        
        
        
        
    def forward(self, x):

        print("input:", x.shape)

        x = self.conv1(x)
        print("conv1:", x.shape)
        x = F.relu(x)
        x = self.pool(x)
        print("pool1:", x.shape)

        x = self.conv2(x)
        print("conv2:", x.shape)
        x = F.relu(x)
        x = self.pool(x)
        print("pool2:", x.shape)

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

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

        x = self.conv5(x)
        print("conv5:", x.shape)
        x = F.relu(x)
        x = self.pool(x)
        print("pool3:", x.shape)

        x = self.flatten(x)
        print("flatten:", x.shape)

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

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

        x = self.linear3(x)
        print("linear3:", x.shape)

        return x
    
        """
        #feature extractor
        print(x.shape)
        out_layer1 = self.pool(F.relu(self.conv1(x)))
        print(out_layer1.shape)
        out_layer2 = self.pool(F.relu(self.conv2(out_layer1)))
        print(out_layer2.shape)
        out_layer3 = F.relu(self.conv3(out_layer2))
        print(out_layer3.shape)
        out_layer4 = F.relu(self.conv4(out_layer3))
        print(out_layer4.shape)
        out_layer5 = self.pool(F.relu(self.conv5(out_layer4))) 
        print(out_layer5.shape)
       
        out_flatten = self.flatten(out_layer5)
        
        # classifier
        out_l1 = F.relu(self.linear1(out_flatten))
        out_l2 = F.relu(self.linear2(out_l1))
        out_l3 = self.linear3(out_l2)
        return out_l3
        """


device = "cpu"

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

print(my_alex)
summary(my_alex, input_size=(1, 3, 224, 224), col_names=("input_size", "output_size", "num_params"), device=device)

"""
# https://pypi.org/project/torchview/
# pip install torchview
from torchview import draw_graph

# https://github.com/mert-kurttutan/torchview
# https://github.com/mert-kurttutan/torchview/issues/113

# graph_dir 'LR' 'TB' 'RL' 'BT'

architecture = "my-alex"
model_graph = draw_graph(my_alex, input_size=(1,3,224,224), graph_dir ='TB' , roll=False, expand_nested=True, graph_name=f'{architecture}',save_graph=False,filename=f'{architecture}')
model_graph.resize_graph(scale=6.0)   
model_graph.visual_graph.render(architecture,format='png')
"""


"""
# https://pytorch.org/vision/main/models/generated/torchvision.models.alexnet.html
torch_alex = torchvision.models.alexnet(weights=None)
print(torch_alex)
#torch_alex.classifier[6] = nn.Linear(4096, 10)
#print(torch_alex)
summary(torch_alex, input_size=(1, 3, 224, 224), col_names=("input_size", "output_size", "num_params"), device=device)
"""