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

from model.base.module import ResidualBlock, ReverseConv2dBlock, Conv2dBlock


class StyleEncoder(nn.Module):
    def __init__(self, in_channels, style_dim, num_conv, end_size=(4, 4), base_channels=64,
                 norm_type="IN", padding_mode='reflect', activation_type="LeakyReLU"):
        super().__init__()
        sequence = [Conv2dBlock(
            in_channels, base_channels, kernel_size=3, stride=1, padding=1, padding_mode=padding_mode,
            activation_type=activation_type, norm_type=norm_type
        )]
        multiple_now = 0
        max_multiple = 3
        for i in range(1, num_conv + 1):
            multiple_prev = multiple_now
            multiple_now = min(2 ** i, 2 ** max_multiple)
            sequence.append(Conv2dBlock(
                multiple_prev * base_channels, multiple_now * base_channels,
                kernel_size=3, stride=2, padding=1, padding_mode=padding_mode,
                activation_type=activation_type, norm_type=norm_type
            ))
        self.sequence = nn.Sequential(*sequence)
        self.fc_avg = nn.Linear(base_channels * (2 ** max_multiple) * end_size[0] * end_size[1], style_dim)
        self.fc_var = nn.Linear(base_channels * (2 ** max_multiple) * end_size[0] * end_size[1], style_dim)

    def forward(self, x):
        x = self.sequence(x)
        x = x.view(x.size(0), -1)
        return self.fc_avg(x), self.fc_var(x)


class SPADEGenerator(nn.Module):
    def __init__(self, in_channels, out_channels, num_blocks, use_vae, num_z_dim, start_size=(4, 4), base_channels=64,
                 padding_mode='reflect', activation_type="LeakyReLU"):
        super().__init__()
        self.sx, self.sy = start_size
        self.use_vae = use_vae
        self.num_z_dim = num_z_dim
        if use_vae:
            self.input_converter = nn.Linear(num_z_dim, 16 * base_channels * self.sx * self.sy)
        else:
            self.input_converter = nn.Conv2d(in_channels, 16 * base_channels, kernel_size=3, padding=1)

        sequence = []

        multiple_now = 16
        for i in range(num_blocks - 1, -1, -1):
            multiple_prev = multiple_now
            multiple_now = min(2 ** i, 2 ** 4)
            if i != num_blocks - 1:
                sequence.append(nn.Upsample(scale_factor=2))
            sequence.append(ResidualBlock(
                base_channels * multiple_prev,
                out_channels=base_channels * multiple_now,
                padding_mode=padding_mode,
                activation_type=activation_type,
                norm_type="SPADE",
                pre_activation=True,
                additional_norm_kwargs=dict(
                    condition_in_channels=in_channels, base_channels=128, base_norm_type="BN",
                    activation_type="ReLU", padding_mode="zeros", gamma_bias=1.0
                )
            ))
        self.sequence = nn.Sequential(*sequence)
        self.output_converter = nn.Sequential(
            ReverseConv2dBlock(base_channels, out_channels, kernel_size=3, stride=1, padding=1,
                               padding_mode=padding_mode, activation_type=activation_type, norm_type="NONE"),
            nn.Tanh()
        )

    def forward(self, seg, z=None):
        if self.use_vae:
            if z is None:
                z = torch.randn(seg.size(0), self.num_z_dim, device=seg.device)
            x = self.input_converter(z).view(seg.size(0), -1, self.sx, self.sy)
        else:
            x = self.input_converter(F.interpolate(seg, size=(self.sx, self.sy)))
        for blk in self.sequence:
            if isinstance(blk, ResidualBlock):
                downsampling_seg = F.interpolate(seg, size=x.size()[2:], mode='nearest')
                blk.conv1.normalization.set_condition_image(downsampling_seg)
                blk.conv2.normalization.set_condition_image(downsampling_seg)
                if blk.learn_skip_connection:
                    blk.res_conv.normalization.set_condition_image(downsampling_seg)
            x = blk(x)
        return self.output_converter(x)

if __name__ == '__main__':
    g = SPADEGenerator(3, 3, 7, False, 256)
    print(g)
    print(g(torch.randn(2, 3, 256, 256)).size())