import math
import functools
import sys

import torch
import torch.nn as nn

from NTED.base_function import EncoderLayer, DecoderLayer, ToRGB
from NTED.edge_attention_layer import Edge_Attn

class Encoder(nn.Module):
    def __init__(
        self, 
        size, 
        input_dim, 
        channels, 
        num_labels=None, 
        match_kernels=None, 
        blur_kernel=[1, 3, 3, 1], 
        ):
        super().__init__()
        self.first = EncoderLayer(input_dim, channels[size], 1)
        self.convs = nn.ModuleList()

        log_size = int(math.log(size, 2))
        self.log_size = log_size

        in_channel = channels[size]
        for i in range(log_size-1, 3, -1):
            out_channel = channels[2 ** i]
            num_label = num_labels[2 ** i] if num_labels is not None else None
            match_kernel = match_kernels[2 ** i] if match_kernels is not None else None
            use_extraction = num_label and match_kernel
            conv = EncoderLayer(
                in_channel, 
                out_channel, 
                kernel_size=3, 
                downsample=True, 
                blur_kernel=blur_kernel,
                use_extraction=use_extraction,
                num_label=num_label,
                match_kernel=match_kernel
                )

            self.convs.append(conv)
            in_channel = out_channel

    def forward(self, input, recoder=None):
        out = self.first(input)
        for idx, layer in enumerate(self.convs):
            out = layer(out, recoder)
        return out

class Decoder(nn.Module):
    def __init__(
        self,
        size,
        channels,
        num_labels,
        match_kernels,
        blur_kernel=[1, 3, 3, 1],
    ):
        super().__init__()


        self.convs = nn.ModuleList()
        # input at resolution 16*16
        in_channel = channels[16]
        self.log_size = int(math.log(size, 2))
        
        for i in range(4, self.log_size + 1):
            out_channel = channels[2 ** i]
            num_label, match_kernel = num_labels[2 ** i], match_kernels[2 ** i]
            use_distribution = num_label and match_kernel
            upsample = (i != 4)
            
            base_layer = functools.partial(
                DecoderLayer,
                out_channel=out_channel,
                kernel_size=3, 
                blur_kernel=blur_kernel,
                use_distribution=use_distribution,
                num_label=num_label,
                match_kernel=match_kernel
                )

            up = nn.Module()   
            up.conv0 = base_layer(in_channel=in_channel, upsample=upsample)
            up.conv1 = base_layer(in_channel=out_channel, upsample=False)
            up.to_rgb = ToRGB(out_channel, upsample=upsample)
            self.convs.append(up)
            in_channel = out_channel 
            
        self.num_labels, self.match_kernels = num_labels, match_kernels
        
        self.edge_attn_block = Edge_Attn(in_channels=3)
    
    def forward(self, input, neural_textures, recoder):
        counter = 0
        out, skip = input, None
        for i, up in enumerate(self.convs):
            if self.num_labels[2**(i+4)] and self.match_kernels[2**(i+4)]:
                neural_texture_conv0 = neural_textures[counter]
                neural_texture_conv1 = neural_textures[counter+1]
                counter += 2
            else:
                neural_texture_conv0, neural_texture_conv1 = None, None
            out = up.conv0(out, neural_texture=neural_texture_conv0, recoder=recoder)
            out = up.conv1(out, neural_texture=neural_texture_conv1, recoder=recoder)
            
            skip = up.to_rgb(out, skip)
        image = self.edge_attn_block(skip)
        # image = skip
        return image