import torch
from torch import nn
import numpy as np
from library.model_layers import ResnetBlock
import math

import torch.nn.functional as F
from torch.nn import init


class Generator(nn.Module):
    def __init__(
        self,
        z_dim=256,
        n_label=10,
        im_size=32,
        im_chan=3,
        embed_size=256,
        nfilter=64,
        nfilter_max=512,
        actvn=nn.ReLU(),
    ):
        super().__init__()
        self.actvn = actvn
        s0 = self.s0 = 4
        nf = self.nf = nfilter
        nf_max = self.nf_max = nfilter_max

        self.z_dim = z_dim
        self.nlabels = n_label
        self.im_size = im_size
        self.im_chan = im_chan

        # Submodules
        nlayers = int(np.log2(im_size / s0))
        self.nf0 = min(nf_max, nf * 2 ** nlayers)

        self.embedding = nn.Embedding(n_label, embed_size)
        self.fc = nn.Linear(z_dim + embed_size, self.nf0 * s0 * s0)

        blocks = []
        for i in range(nlayers):
            nf0 = min(nf * 2 ** (nlayers - i), nf_max)
            nf1 = min(nf * 2 ** (nlayers - i - 1), nf_max)
            blocks += [
                ResnetBlock(nf0, nf1, actvn=self.actvn),
                nn.Upsample(scale_factor=2),
            ]

        blocks += [ResnetBlock(nf, nf, actvn=self.actvn)]

        self.resnet = nn.Sequential(*blocks)
        self.conv_img = nn.Conv2d(nf, im_chan, 3, padding=1)

    def forward(self, z, y):
        assert z.size(0) == y.size(0)
        batch_size = z.size(0)

        yembed = self.embedding(y)
        yembed = yembed / torch.norm(yembed, p=2, dim=1, keepdim=True)
        # print(z.shape, yembed.shape)

        yz = torch.cat([z, yembed], dim=1)
        out = self.fc(yz)
        out = out.view(batch_size, self.nf0, self.s0, self.s0)

        out = self.resnet(out)

        out = self.conv_img(self.actvn(out))
        out = torch.tanh(out)

        return out


class ConditionalBatchNorm2d(nn.BatchNorm2d):

    """Conditional Batch Normalization"""

    def __init__(
        self,
        num_features,
        eps=1e-05,
        momentum=0.1,
        affine=False,
        track_running_stats=True,
    ):
        super(ConditionalBatchNorm2d, self).__init__(
            num_features, eps, momentum, affine, track_running_stats
        )

    def forward(self, input, weight, bias, **kwargs):
        self._check_input_dim(input)

        exponential_average_factor = 0.0

        if self.training and self.track_running_stats:
            self.num_batches_tracked += 1
            if self.momentum is None:  # use cumulative moving average
                exponential_average_factor = 1.0 / self.num_batches_tracked.item()
            else:  # use exponential moving average
                exponential_average_factor = self.momentum

        output = F.batch_norm(
            input,
            self.running_mean,
            self.running_var,
            self.weight,
            self.bias,
            self.training or not self.track_running_stats,
            exponential_average_factor,
            self.eps,
        )
        if weight.dim() == 1:
            weight = weight.unsqueeze(0)
        if bias.dim() == 1:
            bias = bias.unsqueeze(0)
        size = output.size()
        weight = weight.unsqueeze(-1).unsqueeze(-1).expand(size)
        bias = bias.unsqueeze(-1).unsqueeze(-1).expand(size)
        return weight * output + bias


class CategoricalConditionalBatchNorm2d(ConditionalBatchNorm2d):
    def __init__(
        self,
        num_classes,
        num_features,
        eps=1e-5,
        momentum=0.1,
        affine=False,
        track_running_stats=True,
    ):
        super(CategoricalConditionalBatchNorm2d, self).__init__(
            num_features, eps, momentum, affine, track_running_stats
        )
        self.weights = nn.Embedding(num_classes, num_features)
        self.biases = nn.Embedding(num_classes, num_features)

        self._initialize()

    def _initialize(self):
        init.ones_(self.weights.weight.data)
        init.zeros_(self.biases.weight.data)

    def forward(self, input, c, **kwargs):
        weight = self.weights(c)
        bias = self.biases(c)

        return super(CategoricalConditionalBatchNorm2d, self).forward(
            input, weight, bias
        )


def _upsample(x):
    h, w = x.size()[2:]
    return F.interpolate(x, size=(h * 2, w * 2), mode="bilinear")


class Block(nn.Module):
    def __init__(
        self,
        in_ch,
        out_ch,
        h_ch=None,
        ksize=3,
        pad=1,
        activation=F.relu,
        upsample=False,
        num_classes=0,
    ):
        super(Block, self).__init__()

        self.activation = activation
        self.upsample = upsample
        self.learnable_sc = in_ch != out_ch or upsample
        if h_ch is None:
            h_ch = out_ch
        self.num_classes = num_classes

        # Register layrs
        self.c1 = nn.Conv2d(in_ch, h_ch, ksize, 1, pad)
        self.c2 = nn.Conv2d(h_ch, out_ch, ksize, 1, pad)
        if self.num_classes > 0:
            self.b1 = CategoricalConditionalBatchNorm2d(num_classes, in_ch)
            self.b2 = CategoricalConditionalBatchNorm2d(num_classes, h_ch)
        else:
            self.b1 = nn.BatchNorm2d(in_ch)
            self.b2 = nn.BatchNorm2d(h_ch)
        if self.learnable_sc:
            self.c_sc = nn.Conv2d(in_ch, out_ch, 1)

    def _initialize(self):
        init.xavier_uniform_(self.c1.weight.tensor, gain=math.sqrt(2))
        init.xavier_uniform_(self.c2.weight.tensor, gain=math.sqrt(2))
        if self.learnable_sc:
            init.xavier_uniform_(self.c_sc.weight.tensor, gain=1)

    def forward(self, x, y=None, z=None, **kwargs):
        return self.shortcut(x) + self.residual(x, y, z)

    def shortcut(self, x, **kwargs):
        if self.learnable_sc:
            if self.upsample:
                h = _upsample(x)
            h = self.c_sc(h)
            return h
        else:
            return x

    def residual(self, x, y=None, z=None, **kwargs):
        if y is not None:
            h = self.b1(x, y, **kwargs)
        else:
            h = self.b1(x)
        h = self.activation(h)
        if self.upsample:
            h = _upsample(h)
        h = self.c1(h)
        if y is not None:
            h = self.b2(h, y, **kwargs)
        else:
            h = self.b2(h)
        return self.c2(self.activation(h))


class ResNetGenerator(nn.Module):
    """Generator generates 128x128."""

    def __init__(
        self,
        z_dim=256,
        n_label=10,
        im_size=32,
        im_chan=3,
        embed_size=256,
        nfilter=64,
        nfilter_max=512,
        actvn=F.relu,
        distribution="normal",
        bottom_width=4,
    ):
        super(ResNetGenerator, self).__init__()
        self.num_features = num_features = nfilter
        self.dim_z = z_dim
        self.bottom_width = bottom_width
        self.activation = activation = actvn
        self.num_classes = num_classes = n_label
        self.distribution = distribution

        width_coe = 8
        self.l1 = nn.Linear(self.dim_z, width_coe * num_features * bottom_width ** 2)

        self.block2 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block3 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block4 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.b7 = nn.BatchNorm2d(num_features * width_coe)
        self.conv7 = nn.Conv2d(num_features * width_coe, 3, 1, 1)

    def _initialize(self):
        init.xavier_uniform_(self.l1.weight.tensor)
        init.xavier_uniform_(self.conv7.weight.tensor)

    def forward(self, z, y=None, **kwargs):
        h = self.l1(z).view(z.size(0), -1, self.bottom_width, self.bottom_width)
        for i in [2, 3, 4]:
            h = getattr(self, "block{}".format(i))(h, y, **kwargs)
        h = self.activation(self.b7(h))
        return torch.tanh(self.conv7(h))


class ResNetUnconditionalGenerator(nn.Module):
    """Generator generates 128x128."""

    def __init__(
        self,
        z_dim=256,
        n_label=10,
        im_size=32,
        im_chan=3,
        embed_size=256,
        nfilter=64,
        nfilter_max=512,
        actvn=F.relu,
        distribution="normal",
        bottom_width=4,
    ):
        super(ResNetUnconditionalGenerator, self).__init__()
        self.num_features = num_features = nfilter
        self.dim_z = z_dim
        self.bottom_width = bottom_width
        self.activation = activation = actvn
        self.num_classes = num_classes = 0
        self.distribution = distribution

        width_coe = 8
        self.l1 = nn.Linear(self.dim_z, width_coe * num_features * bottom_width ** 2)

        self.block2 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block3 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block4 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.b7 = nn.BatchNorm2d(num_features * width_coe)
        self.conv7 = nn.Conv2d(num_features * width_coe, 3, 1, 1)

    def _initialize(self):
        init.xavier_uniform_(self.l1.weight.tensor)
        init.xavier_uniform_(self.conv7.weight.tensor)

    def forward(self, z, y=None, **kwargs):
        y = None
        h = self.l1(z).view(z.size(0), -1, self.bottom_width, self.bottom_width)
        for i in [2, 3, 4]:
            h = getattr(self, "block{}".format(i))(h, y, **kwargs)
        h = self.activation(self.b7(h))
        return torch.tanh(self.conv7(h))


class ResNetGenerator64(nn.Module):
    """Generator generates 128x128."""

    def __init__(
        self,
        z_dim=256,
        n_label=10,
        im_size=64,
        im_chan=3,
        embed_size=256,
        nfilter=64,
        nfilter_max=512,
        actvn=F.relu,
        distribution="normal",
        bottom_width=4,
    ):
        super(ResNetGenerator64, self).__init__()
        self.num_features = num_features = nfilter
        self.dim_z = z_dim
        self.bottom_width = bottom_width
        self.activation = activation = actvn
        self.num_classes = num_classes = n_label
        self.distribution = distribution

        width_coe = 8
        self.l1 = nn.Linear(self.dim_z, width_coe * num_features * bottom_width ** 2)
        self.block1 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block2 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block3 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block4 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.b7 = nn.BatchNorm2d(num_features * width_coe)
        self.conv7 = nn.Conv2d(num_features * width_coe, 3, 1, 1)

    def _initialize(self):
        init.xavier_uniform_(self.l1.weight.tensor)
        init.xavier_uniform_(self.conv7.weight.tensor)

    def forward(self, z, y=None, **kwargs):
        h = self.l1(z).view(z.size(0), -1, self.bottom_width, self.bottom_width)
        for i in [1, 2, 3, 4]:
            h = getattr(self, "block{}".format(i))(h, y, **kwargs)
        h = self.activation(self.b7(h))
        return torch.tanh(self.conv7(h))


class ResNetGenerator96(nn.Module):
    """Generator generates 128x128."""

    def __init__(
        self,
        z_dim=256,
        n_label=10,
        im_size=32,
        im_chan=3,
        embed_size=256,
        nfilter=64,
        nfilter_max=512,
        actvn=F.relu,
        distribution="normal",
        bottom_width=6,
    ):
        super(ResNetGenerator96, self).__init__()
        self.num_features = num_features = nfilter
        self.dim_z = z_dim
        self.bottom_width = bottom_width
        self.activation = activation = actvn
        self.num_classes = num_classes = n_label
        self.distribution = distribution

        width_coe = 1
        self.l1 = nn.Linear(
            self.dim_z, 16 * width_coe * num_features * bottom_width ** 2
        )

        self.block2 = Block(
            num_features * width_coe * 16,
            num_features * width_coe * 8,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block3 = Block(
            num_features * width_coe * 8,
            num_features * width_coe * 4,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block4 = Block(
            num_features * width_coe * 4,
            num_features * width_coe * 2,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block5 = Block(
            num_features * width_coe * 2,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.b7 = nn.BatchNorm2d(num_features * width_coe)
        self.conv7 = nn.Conv2d(num_features * width_coe, 3, 1, 1)

    def _initialize(self):
        init.xavier_uniform_(self.l1.weight.tensor)
        init.xavier_uniform_(self.conv7.weight.tensor)

    def forward(self, z, y=None, **kwargs):
        h = self.l1(z).view(z.size(0), -1, self.bottom_width, self.bottom_width)
        for i in [2, 3, 4, 5]:
            h = getattr(self, "block{}".format(i))(h, y, **kwargs)
        h = self.activation(self.b7(h))
        return torch.tanh(self.conv7(h))


class ResNetGenerator96_LI(nn.Module):
    """Generator generates 128x128."""

    def __init__(
        self,
        z_dim=256,
        n_label=10,
        im_size=32,
        im_chan=3,
        embed_size=256,
        nfilter=64,
        nfilter_max=512,
        actvn=F.relu,
        distribution="normal",
        bottom_width=6,
    ):
        super(ResNetGenerator96_LI, self).__init__()
        self.num_features = num_features = nfilter
        self.dim_z = z_dim
        self.bottom_width = bottom_width
        self.activation = activation = actvn
        self.num_classes = num_classes = n_label
        self.distribution = distribution

        width_coe = 8
        self.l1 = nn.Linear(
            self.dim_z, 1 * width_coe * num_features * bottom_width ** 2
        )

        self.block2 = Block(
            num_features * width_coe * 1,
            num_features * width_coe * 1,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block3 = Block(
            num_features * width_coe * 1,
            num_features * width_coe * 1,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block4 = Block(
            num_features * width_coe * 1,
            num_features * width_coe * 1,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block5 = Block(
            num_features * width_coe * 1,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.b7 = nn.BatchNorm2d(num_features * width_coe)
        self.conv7 = nn.Conv2d(num_features * width_coe, 3, 1, 1)

    def _initialize(self):
        init.xavier_uniform_(self.l1.weight.tensor)
        init.xavier_uniform_(self.conv7.weight.tensor)

    def forward(self, z, y=None, **kwargs):
        h = self.l1(z).view(z.size(0), -1, self.bottom_width, self.bottom_width)
        for i in [2, 3, 4, 5]:
            h = getattr(self, "block{}".format(i))(h, y, **kwargs)
        h = self.activation(self.b7(h))
        return torch.tanh(self.conv7(h))



generator_dict = {
    "resnet_reggan": Generator,
    "resnet_sngan": ResNetGenerator,
    "resnet_sngan_un": ResNetUnconditionalGenerator,
    "resnet_sngan96": ResNetGenerator96,
    "resnet_sngan64": ResNetGenerator64,
    "resnet_sngan96_li": ResNetGenerator96_LI,
}


import torch
from torch import nn
import numpy as np
from library.model_layers import ResnetBlock
import math

import torch.nn.functional as F
from torch.nn import init


class Generator(nn.Module):
    def __init__(
        self,
        z_dim=256,
        n_label=10,
        im_size=32,
        im_chan=3,
        embed_size=256,
        nfilter=64,
        nfilter_max=512,
        actvn=nn.ReLU(),
    ):
        super().__init__()
        self.actvn = actvn
        s0 = self.s0 = 4
        nf = self.nf = nfilter
        nf_max = self.nf_max = nfilter_max

        self.z_dim = z_dim
        self.nlabels = n_label
        self.im_size = im_size
        self.im_chan = im_chan

        # Submodules
        nlayers = int(np.log2(im_size / s0))
        self.nf0 = min(nf_max, nf * 2 ** nlayers)

        self.embedding = nn.Embedding(n_label, embed_size)
        self.fc = nn.Linear(z_dim + embed_size, self.nf0 * s0 * s0)

        blocks = []
        for i in range(nlayers):
            nf0 = min(nf * 2 ** (nlayers - i), nf_max)
            nf1 = min(nf * 2 ** (nlayers - i - 1), nf_max)
            blocks += [
                ResnetBlock(nf0, nf1, actvn=self.actvn),
                nn.Upsample(scale_factor=2),
            ]

        blocks += [ResnetBlock(nf, nf, actvn=self.actvn)]

        self.resnet = nn.Sequential(*blocks)
        self.conv_img = nn.Conv2d(nf, im_chan, 3, padding=1)

    def forward(self, z, y):
        assert z.size(0) == y.size(0)
        batch_size = z.size(0)

        yembed = self.embedding(y)
        yembed = yembed / torch.norm(yembed, p=2, dim=1, keepdim=True)
        # print(z.shape, yembed.shape)

        yz = torch.cat([z, yembed], dim=1)
        out = self.fc(yz)
        out = out.view(batch_size, self.nf0, self.s0, self.s0)

        out = self.resnet(out)

        out = self.conv_img(self.actvn(out))
        out = torch.tanh(out)

        return out


class ConditionalBatchNorm2d(nn.BatchNorm2d):

    """Conditional Batch Normalization"""

    def __init__(
        self,
        num_features,
        eps=1e-05,
        momentum=0.1,
        affine=False,
        track_running_stats=True,
    ):
        super(ConditionalBatchNorm2d, self).__init__(
            num_features, eps, momentum, affine, track_running_stats
        )

    def forward(self, input, weight, bias, **kwargs):
        self._check_input_dim(input)

        exponential_average_factor = 0.0

        if self.training and self.track_running_stats:
            self.num_batches_tracked += 1
            if self.momentum is None:  # use cumulative moving average
                exponential_average_factor = 1.0 / self.num_batches_tracked.item()
            else:  # use exponential moving average
                exponential_average_factor = self.momentum

        output = F.batch_norm(
            input,
            self.running_mean,
            self.running_var,
            self.weight,
            self.bias,
            self.training or not self.track_running_stats,
            exponential_average_factor,
            self.eps,
        )
        if weight.dim() == 1:
            weight = weight.unsqueeze(0)
        if bias.dim() == 1:
            bias = bias.unsqueeze(0)
        size = output.size()
        weight = weight.unsqueeze(-1).unsqueeze(-1).expand(size)
        bias = bias.unsqueeze(-1).unsqueeze(-1).expand(size)
        return weight * output + bias


class CategoricalConditionalBatchNorm2d(ConditionalBatchNorm2d):
    def __init__(
        self,
        num_classes,
        num_features,
        eps=1e-5,
        momentum=0.1,
        affine=False,
        track_running_stats=True,
    ):
        super(CategoricalConditionalBatchNorm2d, self).__init__(
            num_features, eps, momentum, affine, track_running_stats
        )
        self.weights = nn.Embedding(num_classes, num_features)
        self.biases = nn.Embedding(num_classes, num_features)

        self._initialize()

    def _initialize(self):
        init.ones_(self.weights.weight.data)
        init.zeros_(self.biases.weight.data)

    def forward(self, input, c, **kwargs):
        weight = self.weights(c)
        bias = self.biases(c)

        return super(CategoricalConditionalBatchNorm2d, self).forward(
            input, weight, bias
        )


def _upsample(x):
    h, w = x.size()[2:]
    return F.interpolate(x, size=(h * 2, w * 2), mode="bilinear")


class Block(nn.Module):
    def __init__(
        self,
        in_ch,
        out_ch,
        h_ch=None,
        ksize=3,
        pad=1,
        activation=F.relu,
        upsample=False,
        num_classes=0,
    ):
        super(Block, self).__init__()

        self.activation = activation
        self.upsample = upsample
        self.learnable_sc = in_ch != out_ch or upsample
        if h_ch is None:
            h_ch = out_ch
        self.num_classes = num_classes

        # Register layrs
        self.c1 = nn.Conv2d(in_ch, h_ch, ksize, 1, pad)
        self.c2 = nn.Conv2d(h_ch, out_ch, ksize, 1, pad)
        if self.num_classes > 0:
            self.b1 = CategoricalConditionalBatchNorm2d(num_classes, in_ch)
            self.b2 = CategoricalConditionalBatchNorm2d(num_classes, h_ch)
        else:
            self.b1 = nn.BatchNorm2d(in_ch)
            self.b2 = nn.BatchNorm2d(h_ch)
        if self.learnable_sc:
            self.c_sc = nn.Conv2d(in_ch, out_ch, 1)

    def _initialize(self):
        init.xavier_uniform_(self.c1.weight.tensor, gain=math.sqrt(2))
        init.xavier_uniform_(self.c2.weight.tensor, gain=math.sqrt(2))
        if self.learnable_sc:
            init.xavier_uniform_(self.c_sc.weight.tensor, gain=1)

    def forward(self, x, y=None, z=None, **kwargs):
        return self.shortcut(x) + self.residual(x, y, z)

    def shortcut(self, x, **kwargs):
        if self.learnable_sc:
            if self.upsample:
                h = _upsample(x)
            h = self.c_sc(h)
            return h
        else:
            return x

    def residual(self, x, y=None, z=None, **kwargs):
        if y is not None:
            h = self.b1(x, y, **kwargs)
        else:
            h = self.b1(x)
        h = self.activation(h)
        if self.upsample:
            h = _upsample(h)
        h = self.c1(h)
        if y is not None:
            h = self.b2(h, y, **kwargs)
        else:
            h = self.b2(h)
        return self.c2(self.activation(h))


class ResNetGenerator(nn.Module):
    """Generator generates 128x128."""

    def __init__(
        self,
        z_dim=256,
        n_label=10,
        im_size=32,
        im_chan=3,
        embed_size=256,
        nfilter=64,
        nfilter_max=512,
        actvn=F.relu,
        distribution="normal",
        bottom_width=4,
    ):
        super(ResNetGenerator, self).__init__()
        self.num_features = num_features = nfilter
        self.dim_z = z_dim
        self.bottom_width = bottom_width
        self.activation = activation = actvn
        self.num_classes = num_classes = n_label
        self.distribution = distribution

        width_coe = 8
        self.l1 = nn.Linear(self.dim_z, width_coe * num_features * bottom_width ** 2)

        self.block2 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block3 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block4 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.b7 = nn.BatchNorm2d(num_features * width_coe)
        self.conv7 = nn.Conv2d(num_features * width_coe, 3, 1, 1)

    def _initialize(self):
        init.xavier_uniform_(self.l1.weight.tensor)
        init.xavier_uniform_(self.conv7.weight.tensor)

    def forward(self, z, y=None, **kwargs):
        h = self.l1(z).view(z.size(0), -1, self.bottom_width, self.bottom_width)
        for i in [2, 3, 4]:
            h = getattr(self, "block{}".format(i))(h, y, **kwargs)
        h = self.activation(self.b7(h))
        return torch.tanh(self.conv7(h))


class ResNetUnconditionalGenerator(nn.Module):
    """Generator generates 128x128."""

    def __init__(
        self,
        z_dim=256,
        n_label=10,
        im_size=32,
        im_chan=3,
        embed_size=256,
        nfilter=64,
        nfilter_max=512,
        actvn=F.relu,
        distribution="normal",
        bottom_width=4,
    ):
        super(ResNetUnconditionalGenerator, self).__init__()
        self.num_features = num_features = nfilter
        self.dim_z = z_dim
        self.bottom_width = bottom_width
        self.activation = activation = actvn
        self.num_classes = num_classes = 0
        self.distribution = distribution

        width_coe = 8
        self.l1 = nn.Linear(self.dim_z, width_coe * num_features * bottom_width ** 2)

        self.block2 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block3 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block4 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.b7 = nn.BatchNorm2d(num_features * width_coe)
        self.conv7 = nn.Conv2d(num_features * width_coe, 3, 1, 1)

    def _initialize(self):
        init.xavier_uniform_(self.l1.weight.tensor)
        init.xavier_uniform_(self.conv7.weight.tensor)

    def forward(self, z, y=None, **kwargs):
        y = None
        h = self.l1(z).view(z.size(0), -1, self.bottom_width, self.bottom_width)
        for i in [2, 3, 4]:
            h = getattr(self, "block{}".format(i))(h, y, **kwargs)
        h = self.activation(self.b7(h))
        return torch.tanh(self.conv7(h))


class ResNetGenerator64(nn.Module):
    """Generator generates 128x128."""

    def __init__(
        self,
        z_dim=256,
        n_label=10,
        im_size=64,
        im_chan=3,
        embed_size=256,
        nfilter=64,
        nfilter_max=512,
        actvn=F.relu,
        distribution="normal",
        bottom_width=4,
    ):
        super(ResNetGenerator64, self).__init__()
        self.num_features = num_features = nfilter
        self.dim_z = z_dim
        self.bottom_width = bottom_width
        self.activation = activation = actvn
        self.num_classes = num_classes = n_label
        self.distribution = distribution

        width_coe = 8
        self.l1 = nn.Linear(self.dim_z, width_coe * num_features * bottom_width ** 2)
        self.block1 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block2 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block3 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block4 = Block(
            num_features * width_coe,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.b7 = nn.BatchNorm2d(num_features * width_coe)
        self.conv7 = nn.Conv2d(num_features * width_coe, 3, 1, 1)

    def _initialize(self):
        init.xavier_uniform_(self.l1.weight.tensor)
        init.xavier_uniform_(self.conv7.weight.tensor)

    def forward(self, z, y=None, **kwargs):
        h = self.l1(z).view(z.size(0), -1, self.bottom_width, self.bottom_width)
        for i in [1, 2, 3, 4]:
            h = getattr(self, "block{}".format(i))(h, y, **kwargs)
        h = self.activation(self.b7(h))
        return torch.tanh(self.conv7(h))


class ResNetGenerator96(nn.Module):
    """Generator generates 128x128."""

    def __init__(
        self,
        z_dim=256,
        n_label=10,
        im_size=32,
        im_chan=3,
        embed_size=256,
        nfilter=64,
        nfilter_max=512,
        actvn=F.relu,
        distribution="normal",
        bottom_width=6,
    ):
        super(ResNetGenerator96, self).__init__()
        self.num_features = num_features = nfilter
        self.dim_z = z_dim
        self.bottom_width = bottom_width
        self.activation = activation = actvn
        self.num_classes = num_classes = n_label
        self.distribution = distribution

        width_coe = 1
        self.l1 = nn.Linear(
            self.dim_z, 16 * width_coe * num_features * bottom_width ** 2
        )

        self.block2 = Block(
            num_features * width_coe * 16,
            num_features * width_coe * 8,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block3 = Block(
            num_features * width_coe * 8,
            num_features * width_coe * 4,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block4 = Block(
            num_features * width_coe * 4,
            num_features * width_coe * 2,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block5 = Block(
            num_features * width_coe * 2,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.b7 = nn.BatchNorm2d(num_features * width_coe)
        self.conv7 = nn.Conv2d(num_features * width_coe, 3, 1, 1)

    def _initialize(self):
        init.xavier_uniform_(self.l1.weight.tensor)
        init.xavier_uniform_(self.conv7.weight.tensor)

    def forward(self, z, y=None, **kwargs):
        h = self.l1(z).view(z.size(0), -1, self.bottom_width, self.bottom_width)
        for i in [2, 3, 4, 5]:
            h = getattr(self, "block{}".format(i))(h, y, **kwargs)
        h = self.activation(self.b7(h))
        return torch.tanh(self.conv7(h))


class ResNetGenerator96_LI(nn.Module):
    """Generator generates 128x128."""

    def __init__(
        self,
        z_dim=256,
        n_label=10,
        im_size=32,
        im_chan=3,
        embed_size=256,
        nfilter=64,
        nfilter_max=512,
        actvn=F.relu,
        distribution="normal",
        bottom_width=6,
    ):
        super(ResNetGenerator96_LI, self).__init__()
        self.num_features = num_features = nfilter
        self.dim_z = z_dim
        self.bottom_width = bottom_width
        self.activation = activation = actvn
        self.num_classes = num_classes = n_label
        self.distribution = distribution

        width_coe = 8
        self.l1 = nn.Linear(
            self.dim_z, 1 * width_coe * num_features * bottom_width ** 2
        )

        self.block2 = Block(
            num_features * width_coe * 1,
            num_features * width_coe * 1,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block3 = Block(
            num_features * width_coe * 1,
            num_features * width_coe * 1,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block4 = Block(
            num_features * width_coe * 1,
            num_features * width_coe * 1,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.block5 = Block(
            num_features * width_coe * 1,
            num_features * width_coe,
            activation=activation,
            upsample=True,
            num_classes=num_classes,
        )
        self.b7 = nn.BatchNorm2d(num_features * width_coe)
        self.conv7 = nn.Conv2d(num_features * width_coe, 3, 1, 1)

    def _initialize(self):
        init.xavier_uniform_(self.l1.weight.tensor)
        init.xavier_uniform_(self.conv7.weight.tensor)

    def forward(self, z, y=None, **kwargs):
        h = self.l1(z).view(z.size(0), -1, self.bottom_width, self.bottom_width)
        for i in [2, 3, 4, 5]:
            h = getattr(self, "block{}".format(i))(h, y, **kwargs)
        h = self.activation(self.b7(h))
        return torch.tanh(self.conv7(h))



generator_dict = {
    "resnet_reggan": Generator,
    "resnet_sngan": ResNetGenerator,
    "resnet_sngan_un": ResNetUnconditionalGenerator,
    "resnet_sngan96": ResNetGenerator96,
    "resnet_sngan64": ResNetGenerator64,
    "resnet_sngan96_li": ResNetGenerator96_LI,
}


# Submodules # print(z.shape, yembed.shape) Conditional Batch Normalization # use cumulative moving average # use exponential moving average # Register layrs Generator generates 128x128. Generator generates 128x128. Generator generates 128x128. Generator generates 128x128. Generator generates 128x128.