#!/usr/bin/env python
# coding: utf-8
#
# Author:   Kazuto Nakashima
# URL:      http://kazuto1011.github.io
# Created:  2017-11-19

from __future__ import absolute_import, print_function

import torch
import torch.nn as nn
import torch.nn.functional as F

from .resnet import ConvBlock, PoolBlock, ResnetBlock,ResnetLayer , ResNet


class AtrousSpatialPyramidPooling(nn.Module):
    """
     (ASPP)
    """

    def __init__(self, in_channels, out_channels, rates):
        super(AtrousSpatialPyramidPooling, self).__init__()
        self.conv_layers = nn.ModuleList()
        for i, rate in enumerate(rates):
            conv_layer = nn.Conv2d(in_channels, out_channels, 3, 1, padding=rate, dilation=rate, bias=True)
            torch.nn.init.normal_(conv_layer.weight , 0 , 0.01)
            torch.nn.init.normal_(conv_layer.bias , 0 , 0.01)
            self.conv_layers.append(conv_layer)

    def forward(self, x):
        out = torch.cat([layer(x) for layer in self.conv_layers],dim=1)
        return out

class DeepLabV2(nn.Module):
    def __init__(self, n_classes, n_blocks,atrous_rates):
        super(DeepLabV2, self).__init__()
        ch = [32 * 2 ** p for p in range(6)]
        self.layer1= PoolBlock(ch[0])
        self.layer2= ResnetLayer(n_blocks[0], ch[0], ch[2], 1, 1)
        self.layer3= ResnetLayer(n_blocks[1], ch[2], ch[3], 2, 1)
        self.layer4= ResnetLayer(n_blocks[2], ch[3], ch[4], 2, 1)
        self.atrous = AtrousSpatialPyramidPooling(ch[4], n_classes, atrous_rates)
        self.up_sample = nn.Upsample(scale_factor=2)
        self.conv2d_transpose1 = nn.ConvTranspose2d(in_channels=8,out_channels=4,kernel_size=3,stride=2)
        self.conv2d_transpose2 = nn.ConvTranspose2d(in_channels=4,out_channels=4,kernel_size=3,stride=2)
        self.conv2d = nn.Conv2d(in_channels=4,out_channels=2,kernel_size=3,stride=1)
    
    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.atrous(x)
        x = self.up_sample(x)
        x = self.conv2d_transpose1(x)
        x = self.conv2d_transpose2(x)
        x = self.conv2d(x)
        return x
    
    def freeze_bn(self):
        for m in self.modules():
            if isinstance(m, ConvBlock.BATCH_NORM):
                m.eval()
    
    

# if __name__ == "__main__":
#     model = DeepLabV2(
#         n_classes=21, n_blocks=[3, 4, 23, 3], atrous_rates=[6, 12, 18, 24]
#     )
#     model.eval()
#     image = torch.randn(1, 3, 513, 513)

#     print(model)
#     print("input:", image.shape)
#     print("output:", model(image).shape)