import torch
import torch.nn as nn
import torch.nn.functional as F
import functools
try:
    from .arch_util import EBlock, Attention_Light
    from .arch_util_freq import EBlock_freq
except:
    from arch_util import EBlock, Attention_Light
    from arch_util_freq import EBlock_freq


class Network(nn.Module):
    
    def __init__(self, img_channel=3, 
                 width=16, 
                 middle_blk_num=1, 
                 enc_blk_nums=[], 
                 dec_blk_nums=[],  
                 dilations = [1], 
                 extra_depth_wise = False):
        super(Network, self).__init__()
        
        self.intro = nn.Conv2d(in_channels=img_channel, out_channels=width, kernel_size=3, padding=1, stride=1, groups=1,
                                bias=True)
        self.ending = nn.Conv2d(in_channels=width, out_channels=img_channel, kernel_size=3, padding=1, stride=1, groups=1,
                              bias=True)

        self.encoders = nn.ModuleList()
        self.decoders = nn.ModuleList()
        self.middle_blks = nn.ModuleList()
        self.ups = nn.ModuleList()
        self.downs = nn.ModuleList()

        chan = width
        for num in enc_blk_nums:
            self.encoders.append(
                nn.Sequential(
                    *[EBlock(chan, dilations = dilations, extra_depth_wise=extra_depth_wise) for _ in range(num)]
                )
            )
            self.downs.append(
                nn.Conv2d(chan, 2*chan, 2, 2)
            )
            chan = chan * 2

        self.middle_blks = \
            nn.Sequential(
                *[EBlock(chan, dilations = dilations, extra_depth_wise=extra_depth_wise) for _ in range(middle_blk_num)]
            )

        for num in dec_blk_nums:
            self.ups.append(
                nn.Sequential(
                    nn.Conv2d(chan, chan * 2, 1, bias=False),
                    nn.PixelShuffle(2)
                )
            )
            chan = chan // 2
            self.decoders.append(
                nn.Sequential(
                    *[EBlock(chan, extra_depth_wise=extra_depth_wise) for _ in range(num)]
                )
            )

        self.padder_size = 2 ** len(self.encoders)        
        
        #define the attention layers 
        
        # self.recon_trunk_light = nn.Sequential(*[FBlock(c = chan * self.padder_size,
        #                                         DW_Expand=2, FFN_Expand=2, dilations = dilations, 
        #                                         extra_depth_wise = False) for i in range(residual_layers)])

        # ResidualBlock_noBN_f = functools.partial(ResidualBlock_noBN, nf = width * self.padder_size)
        # self.recon_trunk_light = make_layer(ResidualBlock_noBN_f, residual_layers)
        
   
        
    def forward(self, input):

        _, _, H, W = input.shape

        x = self.intro(input)
        
        encs = []
        # i = 0
        for encoder, down in zip(self.encoders, self.downs):
            x = encoder(x)
            # print(i, x.shape)
            encs.append(x)
            x = down(x)
            # i += 1

        x = self.middle_blks(x)
        # print('3', x.shape)
        # apply the mask
        # x = x * mask
        
        # x = self.recon_trunk_light(x)
        
        for decoder, up, enc_skip in zip(self.decoders, self.ups, encs[::-1]):
            x = up(x)
            x = x + enc_skip
            x = decoder(x)

        x = self.ending(x)
        x = x + input
        
        return x[:, :, :H, :W]


if __name__ == '__main__':
    
    img_channel = 3
    width = 32

    enc_blks = [1, 2, 3]
    middle_blk_num = 3
    dec_blks = [3, 1, 1]
    residual_layers = 2
    dilations = [1, 4]
    
    net = Network(img_channel=img_channel, 
                  width=width, 
                  middle_blk_num=middle_blk_num,
                  enc_blk_nums=enc_blks, 
                  dec_blk_nums=dec_blks,
                  dilations = dilations)

    # NAF = NAFNet(img_channel=img_channel, width=width, middle_blk_num=middle_blk_num,
    #                   enc_blk_nums=enc_blks, dec_blk_nums=dec_blks)

    inp_shape = (3, 256, 256)

    from ptflops import get_model_complexity_info

    macs, params = get_model_complexity_info(net, inp_shape, verbose=False, print_per_layer_stat=False)

    print(macs, params)    
    inp = torch.randn(1, 3, 256, 256)
    out = net(inp)