import numpy as np
import torch
import torch.nn.functional as F
from torch import nn
from torch.nn import Linear, Conv2d, BatchNorm2d, PReLU, Sequential, Module

from models.encoders.helpers import get_blocks, Flatten, bottleneck_IR, bottleneck_IR_SE
from models.stylegan2.model import EqualLinear


class GradualStyleBlock(Module):
    def __init__(self, in_c, out_c, spatial, max_pooling=False):
        super(GradualStyleBlock, self).__init__()
        self.out_c = out_c
        self.spatial = spatial
        self.max_pooling = max_pooling
        num_pools = int(np.log2(spatial))
        modules = []
        modules += [Conv2d(in_c, out_c, kernel_size=3, stride=2, padding=1),
                    nn.LeakyReLU()]
        for i in range(num_pools - 1):
            modules += [
                Conv2d(out_c, out_c, kernel_size=3, stride=2, padding=1),
                nn.LeakyReLU()
            ]
        self.convs = nn.Sequential(*modules)
        self.linear = EqualLinear(out_c, out_c, lr_mul=1)

    def forward(self, x):
        x = self.convs(x)
        # To make E accept more general H*W images, we add global average pooling to 
        # resize all features to 1*1*512 before mapping to latent codes
        if self.max_pooling:
            x = F.adaptive_max_pool2d(x, 1) ##### modified
        else:
            x = F.adaptive_avg_pool2d(x, 1) ##### modified
        x = x.view(-1, self.out_c)
        x = self.linear(x)
        return x

class AdaptiveInstanceNorm(nn.Module):
    def __init__(self, fin, style_dim=512):
        super().__init__()

        self.norm = nn.InstanceNorm2d(fin, affine=False)
        self.style = nn.Linear(style_dim, fin * 2)

        self.style.bias.data[:fin] = 1
        self.style.bias.data[fin:] = 0

    def forward(self, input, style):
        style = self.style(style).unsqueeze(2).unsqueeze(3)
        gamma, beta = style.chunk(2, 1)
        out = self.norm(input)
        out = gamma * out + beta
        return out    
    

class FusionLayer(Module):  ##### modified 
    def __init__(self, inchannel, outchannel, use_skip_torgb=True, use_att=0):
        super(FusionLayer, self).__init__()
        
        self.transform = nn.Sequential(nn.Conv2d(inchannel, outchannel, kernel_size=3, stride=1, padding=1),
                                      nn.LeakyReLU())
        self.fusion_out = nn.Conv2d(outchannel*2, outchannel, kernel_size=3, stride=1, padding=1)
        self.fusion_out.weight.data *= 0.01
        self.fusion_out.weight[:,0:outchannel,1,1].data += torch.eye(outchannel) 
        
        self.use_skip_torgb = use_skip_torgb
        if use_skip_torgb:
            self.fusion_skip = nn.Conv2d(3+outchannel, 3, kernel_size=3, stride=1, padding=1)
            self.fusion_skip.weight.data *= 0.01
            self.fusion_skip.weight[:,0:3,1,1].data += torch.eye(3)
            
        self.use_att = use_att
        if use_att:
            modules = []
            modules.append(nn.Linear(512, outchannel))
            for _ in range(use_att):
                modules.append(nn.LeakyReLU(negative_slope=0.2, inplace=True))
                modules.append(nn.Linear(outchannel, outchannel))
            modules.append(nn.LeakyReLU(negative_slope=0.2, inplace=True))
            self.linear = Sequential(*modules)
            self.norm = AdaptiveInstanceNorm(outchannel*2, outchannel)
            self.conv = nn.Conv2d(outchannel*2, 1, 3, 1, 1, bias=True)

    def forward(self, feat, out, skip, editing_w=None):
        x = self.transform(feat)
        # similar to VToonify, use editing vector as condition
        # fuse encoder feature and decoder feature with a predicted attention mask m_E
        # if self.use_att = False, just fuse them with a simple conv layer
        if self.use_att and editing_w is not None:
            label = self.linear(editing_w)
            m_E = (F.relu(self.conv(self.norm(torch.cat([out, abs(out-x)], dim=1), label)))).tanh()
            x = x * m_E
        out = self.fusion_out(torch.cat((out, x), dim=1))
        if self.use_skip_torgb:
            skip = self.fusion_skip(torch.cat((skip, x), dim=1))
        return out, skip

    
class ResnetBlock(nn.Module):
    def __init__(self, dim):
        super(ResnetBlock, self).__init__()

        self.conv_block = nn.Sequential(Conv2d(dim, dim, 3, 1, 1),
                                        nn.LeakyReLU(),
                                        Conv2d(dim, dim, 3, 1, 1))
        self.relu = nn.LeakyReLU()

    def forward(self, x):
        out = x + self.conv_block(x)  
        return self.relu(out)

# trainable light-weight translation network T
# for sketch/mask-to-face translation, 
# we add a trainable T to map y to an intermediate domain where E can more easily extract features. 
class ResnetGenerator(nn.Module):
    def __init__(self, in_channel=19, res_num=2):
        super(ResnetGenerator, self).__init__()
        
        modules = []
        modules.append(Conv2d(in_channel, 16, 3, 2, 1))
        modules.append(nn.LeakyReLU())
        modules.append(Conv2d(16, 16, 3, 2, 1))
        modules.append(nn.LeakyReLU())
        for _ in range(res_num):
            modules.append(ResnetBlock(16))
        for _ in range(2):
            modules.append(nn.ConvTranspose2d(16, 16, 3, 2, 1, output_padding=1))
            modules.append(nn.LeakyReLU())  
        modules.append(Conv2d(16, 64, 3, 1, 1, bias=False))
        modules.append(BatchNorm2d(64))  
        modules.append(PReLU(64))  
        self.model = Sequential(*modules)

    def forward(self, input):
        return self.model(input)
    
class GradualStyleEncoder(Module):
    def __init__(self, num_layers, mode='ir', opts=None):
        super(GradualStyleEncoder, self).__init__()
        assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152'
        assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
        blocks = get_blocks(num_layers)
        if mode == 'ir':
            unit_module = bottleneck_IR
        elif mode == 'ir_se':
            unit_module = bottleneck_IR_SE
        
        # for sketch/mask-to-face translation, add a new network T
        if opts.input_nc != 3:
            self.input_label_layer = ResnetGenerator(opts.input_nc, opts.res_num)

        self.input_layer = Sequential(Conv2d(3, 64, (3, 3), 1, 1, bias=False),
                                      BatchNorm2d(64),
                                      PReLU(64))
        modules = []
        for block in blocks:
            for bottleneck in block:
                modules.append(unit_module(bottleneck.in_channel,
                                           bottleneck.depth,
                                           bottleneck.stride))
        self.body = Sequential(*modules)

        self.styles = nn.ModuleList()
        self.style_count = opts.n_styles
        self.coarse_ind = 3
        self.middle_ind = 7
        for i in range(self.style_count):
            if i < self.coarse_ind:
                style = GradualStyleBlock(512, 512, 16, 'max_pooling' in opts and opts.max_pooling)
            elif i < self.middle_ind:
                style = GradualStyleBlock(512, 512, 32, 'max_pooling' in opts and opts.max_pooling)
            else:
                style = GradualStyleBlock(512, 512, 64, 'max_pooling' in opts and opts.max_pooling)
            self.styles.append(style)
        self.latlayer1 = nn.Conv2d(256, 512, kernel_size=1, stride=1, padding=0)
        self.latlayer2 = nn.Conv2d(128, 512, kernel_size=1, stride=1, padding=0)
        
        # we concatenate pSp features in the middle layers and 
        # add a convolution layer to map the concatenated features to the first-layer input feature f of G.
        self.featlayer = nn.Conv2d(768, 512, kernel_size=1, stride=1, padding=0) ##### modified
        self.skiplayer = nn.Conv2d(768, 3, kernel_size=1, stride=1, padding=0) ##### modified
        
        # skip connection
        if 'use_skip' in opts and opts.use_skip: ##### modified
            self.fusion = nn.ModuleList()
            channels = [[256,512], [256,512], [256,512], [256,512], [128,512], [64,256], [64,128]]
            # opts.skip_max_layer: how many layers are skipped to the decoder
            for inc, outc in channels[:max(1, min(7, opts.skip_max_layer))]: # from 4 to 256
                self.fusion.append(FusionLayer(inc, outc, opts.use_skip_torgb, opts.use_att))

    def _upsample_add(self, x, y):
        '''Upsample and add two feature maps.
        Args:
          x: (Variable) top feature map to be upsampled.
          y: (Variable) lateral feature map.
        Returns:
          (Variable) added feature map.
        Note in PyTorch, when input size is odd, the upsampled feature map
        with `F.upsample(..., scale_factor=2, mode='nearest')`
        maybe not equal to the lateral feature map size.
        e.g.
        original input size: [N,_,15,15] ->
        conv2d feature map size: [N,_,8,8] ->
        upsampled feature map size: [N,_,16,16]
        So we choose bilinear upsample which supports arbitrary output sizes.
        '''
        _, _, H, W = y.size()
        return F.interpolate(x, size=(H, W), mode='bilinear', align_corners=True) + y
    
    # return_feat: return f
    # return_full: return f and the skipped encoder features
    # return [out, feats]
    # out is the style latent code w+
    # feats[0] is f for the 1st conv layer, feats[1] is f for the 1st torgb layer
    # feats[2-8] is the skipped encoder features 
    def forward(self, x, return_feat=False, return_full=False): ##### modified      
        if x.shape[1] != 3:
            x = self.input_label_layer(x)
        else:
            x = self.input_layer(x)
        c256 = x ##### modified
        
        latents = []
        modulelist = list(self.body._modules.values())
        for i, l in enumerate(modulelist):
            x = l(x)
            if i == 2:    ##### modified
                c128 = x
            elif i == 6:
                c1 = x
            elif i == 10: ##### modified
                c21 = x   ##### modified
            elif i == 15: ##### modified
                c22 = x   ##### modified
            elif i == 20:
                c2 = x
            elif i == 23:
                c3 = x

        for j in range(self.coarse_ind):
            latents.append(self.styles[j](c3))

        p2 = self._upsample_add(c3, self.latlayer1(c2))
        for j in range(self.coarse_ind, self.middle_ind):
            latents.append(self.styles[j](p2))

        p1 = self._upsample_add(p2, self.latlayer2(c1))
        for j in range(self.middle_ind, self.style_count):
            latents.append(self.styles[j](p1))

        out = torch.stack(latents, dim=1)
        
        if not return_feat:
            return out
        
        feats = [self.featlayer(torch.cat((c21, c22, c2), dim=1)), self.skiplayer(torch.cat((c21, c22, c2), dim=1))]
        
        if return_full: ##### modified
            feats += [c2, c2, c22, c21, c1, c128, c256]
            
        return out, feats

    
    # only compute the first-layer feature f
    # E_F in the paper
    def get_feat(self, x): ##### modified     
        # for sketch/mask-to-face translation
        # use a trainable light-weight translation network T
        if x.shape[1] != 3:
            x = self.input_label_layer(x)
        else:
            x = self.input_layer(x)
        
        latents = []
        modulelist = list(self.body._modules.values())
        for i, l in enumerate(modulelist):
            x = l(x)
            if i == 10: ##### modified
                c21 = x   ##### modified
            elif i == 15: ##### modified
                c22 = x   ##### modified
            elif i == 20:
                c2 = x
                break
        return self.featlayer(torch.cat((c21, c22, c2), dim=1))
    
class BackboneEncoderUsingLastLayerIntoW(Module):
    def __init__(self, num_layers, mode='ir', opts=None):
        super(BackboneEncoderUsingLastLayerIntoW, self).__init__()
        print('Using BackboneEncoderUsingLastLayerIntoW')
        assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152'
        assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
        blocks = get_blocks(num_layers)
        if mode == 'ir':
            unit_module = bottleneck_IR
        elif mode == 'ir_se':
            unit_module = bottleneck_IR_SE
        self.input_layer = Sequential(Conv2d(opts.input_nc, 64, (3, 3), 1, 1, bias=False),
                                      BatchNorm2d(64),
                                      PReLU(64))
        self.output_pool = torch.nn.AdaptiveAvgPool2d((1, 1))
        self.linear = EqualLinear(512, 512, lr_mul=1)
        modules = []
        for block in blocks:
            for bottleneck in block:
                modules.append(unit_module(bottleneck.in_channel,
                                           bottleneck.depth,
                                           bottleneck.stride))
        self.body = Sequential(*modules)

    def forward(self, x):
        x = self.input_layer(x)
        x = self.body(x)
        x = self.output_pool(x)
        x = x.view(-1, 512)
        x = self.linear(x)
        return x


class BackboneEncoderUsingLastLayerIntoWPlus(Module):
    def __init__(self, num_layers, mode='ir', opts=None):
        super(BackboneEncoderUsingLastLayerIntoWPlus, self).__init__()
        print('Using BackboneEncoderUsingLastLayerIntoWPlus')
        assert num_layers in [50, 100, 152], 'num_layers should be 50,100, or 152'
        assert mode in ['ir', 'ir_se'], 'mode should be ir or ir_se'
        blocks = get_blocks(num_layers)
        if mode == 'ir':
            unit_module = bottleneck_IR
        elif mode == 'ir_se':
            unit_module = bottleneck_IR_SE
        self.n_styles = opts.n_styles
        self.input_layer = Sequential(Conv2d(opts.input_nc, 64, (3, 3), 1, 1, bias=False),
                                      BatchNorm2d(64),
                                      PReLU(64))
        self.output_layer_2 = Sequential(BatchNorm2d(512),
                                         torch.nn.AdaptiveAvgPool2d((7, 7)),
                                         Flatten(),
                                         Linear(512 * 7 * 7, 512))
        self.linear = EqualLinear(512, 512 * self.n_styles, lr_mul=1)
        modules = []
        for block in blocks:
            for bottleneck in block:
                modules.append(unit_module(bottleneck.in_channel,
                                           bottleneck.depth,
                                           bottleneck.stride))
        self.body = Sequential(*modules)

    def forward(self, x):
        x = self.input_layer(x)
        x = self.body(x)
        x = self.output_layer_2(x)
        x = self.linear(x)
        x = x.view(-1, self.n_styles, 512)
        return x