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

# https://arxiv.org/pdf/1704.04861v1

# ukázka klasické konvoluční vrstva
class StandardConvExample(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=5, padding=0, stride=1, bias=False)

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

#Zde se rozděluje standardní konvoluci na dva lehčí kroky (známé jako Depthwise Separable Convolution): 
# - Depthwise (groups=in_channels): Aplikuje filtr (např. 5x5) na každý kanál zvlášť. Nemíchá kanály mezi sebou. 
# - Pointwise (kernel_size=1): Aplikuje konvoluci 1x1, která se dívá napříč kanály a smíchá je dohromady (změní počet kanálů např. ze 3 na 256).

# - Depthwise část (groups=in_channels) izoluje každý kanál a aplikuje na něj filtr zvlášť.
# - Pointwise část (kernel_size=1) následně "promíchá" informace napříč kanály, aby vytvořila nové rysy (features).
class DepthPointWiseConvExample(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.depthwise = nn.Conv2d(in_channels, in_channels, kernel_size=5, padding=0, groups=in_channels, bias=False)
        self.pointwise = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)

    def forward(self, X):
        #print("X.shape", X.shape)
        out = self.depthwise(X)
        #print("depthwise.shape", out.shape)
        out = self.pointwise(out)
        return out


# Case 1: Standardní konvoluce ("malá")
print("--------------------------------------")
print("------------------- CASE - 1 StandardConv")
print("--------------------------------------")
batch_size = 1
in_channels = 3
img_size = 12
out_channels = 1

my_input = torch.randn(batch_size, in_channels, img_size, img_size)
standard_conv_block = StandardConvExample(in_channels, out_channels)
standard_conv_block_out = standard_conv_block(my_input)

print("standard_conv_block_out.shape", standard_conv_block_out.shape)
summary(standard_conv_block, input_size=(batch_size, in_channels, img_size, img_size), col_names=["kernel_size", "output_size", "num_params", "mult_adds"],)

#standard_conv_block_out.shape torch.Size([1, 1, 8, 8])
#============================================================================================================================================
#Layer (type:depth-idx)                   Kernel Shape              Output Shape              Param #                   Mult-Adds
#============================================================================================================================================
#StandardConvExample                      --                        [1, 1, 8, 8]              --                        --
#├─Conv2d: 1-1                            [5, 5]                    [1, 1, 8, 8]              75                        4,800

# Case 2: Standardní konvoluce ("velká")
print("--------------------------------------")
print("------------------- CASE - 2 StandardConv")
print("--------------------------------------")
batch_size = 1
in_channels = 3
img_size = 12
out_channels = 256

my_input = torch.randn(batch_size, in_channels, img_size, img_size)
standard_conv_block = StandardConvExample(in_channels, out_channels)
standard_conv_block_out = standard_conv_block(my_input)

print("standard_conv_block_out.shape", standard_conv_block_out.shape)
summary(standard_conv_block, input_size=(batch_size, in_channels, img_size, img_size), col_names=["kernel_size", "output_size", "num_params", "mult_adds"],)

#standard_conv_block_out.shape torch.Size([1, 256, 8, 8])
#============================================================================================================================================
#Layer (type:depth-idx)                   Kernel Shape              Output Shape              Param #                   Mult-Adds
#============================================================================================================================================
#StandardConvExample                      --                        [1, 256, 8, 8]            --                        --
#├─Conv2d: 1-1                            [5, 5]                    [1, 256, 8, 8]            19,200                    1,228,800
#============================================================================================================================================

# Case 3: Depthwise Separable konvoluce (MobileNet styl)
# Dosahuje stejného tvaru výstupu jako u klasické konvoluci ([1, 256, 8, 8]), ale mnohem efektivněji.
print("--------------------------------------")
print("------------------- CASE - 3 DepthPointWiseConv")
print("--------------------------------------")
batch_size = 1
in_channels = 3
img_size = 12
out_channels = 256

my_input = torch.randn(batch_size, in_channels, img_size, img_size)
depthpoint_conv_block = DepthPointWiseConvExample(in_channels, out_channels)
depthpoint_conv_block_out = depthpoint_conv_block(my_input)

print("depthpoint_conv_block_out.shape", depthpoint_conv_block_out.shape)
summary(depthpoint_conv_block, input_size=(batch_size, in_channels, img_size, img_size), col_names=["kernel_size", "output_size", "num_params", "mult_adds"],)

#depthpoint_conv_block_out.shape torch.Size([1, 256, 8, 8])
#============================================================================================================================================
#Layer (type:depth-idx)                   Kernel Shape              Output Shape              Param #                   Mult-Adds
#============================================================================================================================================
#DepthPointWiseConvExample                --                        [1, 256, 8, 8]            --                        --
#├─Conv2d: 1-1                            [5, 5]                    [1, 3, 8, 8]              75                        4,800
#├─Conv2d: 1-2                            [1, 1]                    [1, 256, 8, 8]            768                       49,152
#============================================================================================================================================