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

class conv_block(nn.Module):
    def __init__(self, in_c, out_c, act):
        super().__init__()

        self.conv1 = nn.Conv2d(in_c, out_c, kernel_size=3, padding=1)
        self.bn1 = nn.BatchNorm2d(out_c)

        self.conv2 = nn.Conv2d(out_c, out_c, kernel_size=3, padding=1)
        self.bn2 = nn.BatchNorm2d(out_c)

        if act == 'relu':
            self.activation = nn.ReLU()
        elif act == 'sigmoid':
            self.activation = nn.Sigmoid()
        else:
            self.activation = nn.Identity()

        # self.relu = nn.ReLU()

    def forward(self, inputs):
        x = self.conv1(inputs)
        x = self.bn1(x)
        x = self.activation(x)

        x = self.conv2(x)
        x = self.bn2(x)
        x = self.activation(x)

        return x
class Decoder_block(nn.Module):
    def __init__(self, in_channel, out_channel, kernel, stride, padding=1, out_padding=1, act = 'relu') -> None:
        super().__init__()
        self.upsample = nn.ConvTranspose2d(in_channels=in_channel,\
                                            out_channels=out_channel,\
                                            kernel_size=kernel,\
                                             stride=stride,\
                                              padding=padding,
                                              output_padding=out_padding)
        if act == 'relu':
            self.activation = nn.ReLU()
        elif act == 'sigmoid':
            self.activation = nn.Sigmoid()
        else:
            self.activation = nn.Identity()
        
    def forward(self, x):
        return self.activation(self.upsample(x))
    
class Decoder(nn.Module):
    def __init__(self, num_layers, channels, kernels, strides, activations) -> None:
        super().__init__()
        assert len(channels) -1 == len(kernels) and len(strides) == len(kernels) and num_layers == len(strides)
        assert num_layers == len(activations)
        self.layers = []
        for i in range(num_layers):
            self.layers.append(Decoder_block(in_channel=channels[i],\
                                             out_channel=channels[i+1],\
                                              kernel=kernels[i],\
                                                stride=strides[i],\
                                                 act=activations[i]))
            self.layers.append(conv_block(in_c=channels[i+1],out_c=channels[i+1], act= activations[i]))
        self.model = nn.Sequential(*self.layers)
    def forward(self, x):
        return self.model(x)

class enc_dec_model(nn.Module):
    def __init__(self, max_depth=10, backbone='resnet', unfreeze = False) -> None:
        super().__init__()
        if backbone == 'resnet':
            self.encoder = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)
            num_layers=5
            channels=[2048,256,128,64,32,1]
            kernels=[3,3,3,3,3]
            strides = [2,2,2,2,2]
            activations=['relu', 'relu', 'relu' ,'relu', 'sigmoid']
        if unfreeze:
            for param in self.encoder.parameters():
                param.requires_grad = True
        else:    
            for param in self.encoder.parameters():
                param.requires_grad = False
            for i, child in enumerate(self.encoder.children()):
                if i == 7:
                    for j, child2 in enumerate(child.children()):
                        if j == 2:
                            # print("count:", j)
                            # print(child2)
                            for param in child2.parameters():
                                param.requires_grad = True
                if i>=8:
                    # print("count:", i)
                    # print(child)
                    for param in child.parameters():
                        param.requires_grad = True
        # input(":")
        self.encoder = torch.nn.Sequential(*(list(self.encoder.children())[:-2]))
        # self.bridge = nn.Conv2d(2048, 2048, 1, 1)

        self.decoder = Decoder(num_layers=num_layers,\
                                channels=channels,\
                                kernels=kernels,\
                                strides = strides,\
                                 activations=activations)
        self.max_depth = max_depth
    def forward(self, x):
        x = self.encoder(x)
        # x = self.bridge(x)
        # print(x)
        x = self.decoder(x)
        # print(x)
        x = x*self.max_depth
        return {'pred_d':x}
    
if __name__ == "__main__":
    # model = Decoder(num_layers=5,\
    #                 channels=[2048,256,128,64,32,1],\
    #                 kernels=[3,3,3,3,3],\
    #                 strides = [2,2,2,2,2])
    model = enc_dec_model(unfreeze=True).cuda()
    print(model)
    summary(model, input_size=(64,3,448,448))