<reponame>Penchekrak/DeepGenerativeModels
import torchvision
import wandb
from hydra.utils import instantiate
from pytorch_lightning import LightningModule
from torch import nn
import torch
from torch.nn import functional as F
import typing as tp
from utils import FidScore

EPS = 1e-18

class UnNormalize(object):
    def __init__(self, mean, std):
        self.mean = mean
        self.std = std

    def __call__(self, tensor):
        """
        Args:
            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
        Returns:
            Tensor: Normalized image.
        """
        for t, m, s in zip(tensor, self.mean, self.std):
            t.mul_(s).add_(m)
            # The normalize code -> t.sub_(m).div_(s)
        return tensor


def gaussian_log_likelihood(x_true, mu, sigma):
    """
    Compute log-likelihood of objects x_true for the generated by model
    component-wise Gaussian distributions.
    Each object from x_true has K corresponding distrbutions from x_distr.
    Log-likelihood estimation must be computed for each pair of an object
    and a corresponding to the object distribution.
    Do not forget about computational stability!
    Do not divide log-likelihood by the dimensionality of the space of objects.
    Input: x_true, Tensor of shape n x D.
    Input: x_distr, Tensor of shape n x K x D - parameters of component-wise
           Gaussian distributions.
    Return: Tensor of shape n x K - log-likelihood for each pair of an object
            and a corresponding distribution.
    """
    log_pdf = - torch.sum(torch.log(sigma + EPS) + 0.5 * ((x_true.unsqueeze(1) - mu) / sigma) ** 2,
                          dim=(-1, -2, -3))  # - D / 2 * log(2 * pi)
    return log_pdf


def kl(q_mu, q_sigma, p_mu, p_sigma):
    """

    Compute KL-divergence KL(q || p) between n pairs of Gaussians
    with diagonal covariational matrices.
    Do not divide KL-divergence by the dimensionality of the latent space.

    Input: q_mu, p_mu, Tensor of shape n x d - mean vectors for n Gaussians.
    Input: q_sigma, p_sigma, Tensor of shape n x d - standard deviation
           vectors for n Gaussians.
    Return: Tensor of shape n - each component is KL-divergence between
            a corresponding pair of Gaussians.
    """
    return (torch.log(p_sigma + EPS) -
            torch.log(q_sigma + EPS) +
            (q_sigma ** 2 + (q_mu - p_mu) ** 2) / (2 * p_sigma ** 2 + EPS) -
            0.5).sum(-1)


class ClampLayer(nn.Module):
    def __init__(self, min=None, max=None):
        super().__init__()
        self.min = min
        self.max = max
        self.kwargs = {}
        if min is not None:
            self.kwargs['min'] = min
        if max is not None:
            self.kwargs['max'] = max

    def forward(self, input):
        return torch.clamp(input, **self.kwargs)


class Block(nn.Module):
    def __init__(self, in_features, out_features, kernel, stride=1, bias=False, upsample=False):
        super().__init__()
        self.upsample = upsample

        self.conv = nn.Conv2d(in_features, out_features, kernel, stride=stride, padding=(kernel - 1) // 2, bias=bias)
        self.norm = nn.BatchNorm2d(out_features)
        self.act = nn.LeakyReLU(0.2)

    def forward(self, x):
        if self.upsample:
            x = F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False, recompute_scale_factor=False)
        return self.act(self.norm(self.conv(x)))


class VAE(LightningModule):
    def __init__(self, latent_dim, optimizer_conf=None, *args, **kwargs):
        """
        Initialize model weights.
        Input: d, int - the dimensionality of the latent space.
        Input: D, int - the dimensionality of the object space.
        """
        super(VAE, self).__init__(*args, **kwargs)
        self.latent_dim = latent_dim
        self.optimizer_conf = optimizer_conf
        self.proposal_network = nn.Sequential(
            Block(3, 16, 3, stride=2),
            Block(16, 32, 3, stride=2),
            Block(32, 32, 3, stride=2),
            Block(32, 32, 3, stride=2),
            Block(32, 64, 3, stride=1).conv,
            nn.AdaptiveAvgPool3d(output_size=(64, 1, 1)),
            nn.Flatten()
        )

        self.proposal_mu_head = nn.Linear(64, self.latent_dim)
        self.proposal_sigma_head = nn.Sequential(
            nn.Linear(64, self.latent_dim),
            nn.Softplus()
        )

        self.generative_network = nn.Sequential(
            nn.Unflatten(1, (self.latent_dim, 1, 1)),
            nn.Conv2d(self.latent_dim, 32, 2, padding=1),
            nn.BatchNorm2d(32),
            nn.LeakyReLU(0.2),
            Block(32, 32, 3, upsample=True),
            Block(32, 32, 3, upsample=True),
            Block(32, 32, 3, upsample=True),
            Block(32, 32, 3, upsample=True),
            Block(32, 32, 3, upsample=True),
            Block(32, 32, 3).conv,
        )

        self.generative_mu_head = nn.Sequential(
            nn.Conv2d(32, 3, 1),
            ClampLayer(-10, 10),
            nn.Sigmoid()
        )
        self.generative_sigma_head = nn.Sequential(
            nn.Conv2d(32, 3, 1),
            ClampLayer(1e-3, None),
            nn.Softplus()
        )
        self.fid = FidScore()

    def proposal_distr(self, x):
        """
        Generate proposal distribution over z.
        Note that sigma is positive by design of neural network.
        Input: x, Tensor of shape n x D.
        Return: tuple(Tensor, Tensor),
                Each Tensor is a matrix of shape n x d.
                The first one is mu, the second one is sigma.
        """
        proposal = self.proposal_network(x)
        mu = self.proposal_mu_head(proposal)
        sigma = self.proposal_sigma_head(proposal)
        return mu, sigma

    def prior_distr(self, x):
        """
        Generate prior distribution over z.
        Note that sigma is positive by design of neural network.
        Input: x, Tensor of shape n x D.
        Return: tuple(Tensor, Tensor),
                Each Tensor is a matrix of shape n x d.
                The first one is mu, the second one is sigma.
        """
        mu = torch.zeros((x.shape[0], self.latent_dim)).type_as(x)
        sigma = torch.ones((x.shape[0], self.latent_dim)).type_as(x)
        return mu, sigma

    def sample_latent(self, mu, sigma, K=1):
        """
        Generate samples from Gaussians with diagonal covariance matrices in latent space.
        Input: mu, Tensor of shape n x d - mean vectors for n Gaussians.
        Input: sigma, Tensor of shape n x d - standard deviation vectors
               for n Gaussians.
        Input: K, int - number of samples from each Gaussian.
        Return: Tensor of shape n x K x d.
        """
        epsilon = torch.randn((mu.shape[0], K, mu.shape[1])).type_as(mu)
        sample = mu.unsqueeze(1) + sigma.unsqueeze(1) * epsilon
        return sample

    def generative_distr(self, z):
        """
        Compute a tensor of parameters of Bernoulli distribution over x
        given a tensor of latent representations.
        Input: z, Tensor of shape n x K x d - tensor of latent representations.
        Return: Tensor of shape n x K x D - parameters of Bernoulli distribution.
        """
        n, K, d = z.shape
        generates = self.generative_network(z.view((n * K, d)))
        nK, C, H, W = generates.shape
        mu = self.generative_mu_head(generates)
        sigma = self.generative_sigma_head(generates)
        return mu.view((n, K, -1, H, W)), sigma.view((n, K, -1, H, W))

    def training_step(self, batch, batch_idx):
        """
        Compute VLB for batch. The VLB for batch is an average of VLBs for batch's objects.
        Input: batch, Tensor of shape n x D.
        Return: Tensor, scalar - VLB.
        """
        batch, _ = batch
        mu, sigma = self.proposal_distr(batch)
        mu_prior, sigma_prior = self.prior_distr(batch)
        z = self.sample_latent(mu, sigma)
        mu_image, sigma_image = self.generative_distr(z)
        gll = gaussian_log_likelihood(batch, mu_image, sigma_image).mean()
        kl_ = kl(mu, sigma, mu_prior, sigma_prior).mean()
        vlb = gll - kl_
        self.log_dict(
            {
                'gaussian log-likelihood': gll,
                'kl divergence': kl_,
                'total': vlb
            }
        )
        return -vlb

    def forward(self, batch, K=1, *args, **kwargs):
        mu, sigma = self.proposal_distr(batch)
        z = self.sample_latent(mu, sigma, K)
        mu_image, sigma_image = self.generative_distr(z)
        return mu_image, sigma_image

    @torch.no_grad()
    def generate_samples(self, num_samples=1):
        """
        Generate samples from the model.
        Input: num_samples, int - number of samples to generate.
        Return: Tensor of shape num_samples x D.
        """
        device = next(iter(self.generative_network.parameters())).device
        mu, sigma = self.prior_distr(torch.empty((1, self.latent_dim), device=device, dtype=torch.float))
        z = self.sample_latent(mu, sigma, K=num_samples)
        samples, _ = self.generative_distr(z)
        return samples.squeeze(0)

    def validation_step(
            self,
            batch: tp.Tuple[torch.Tensor, torch.Tensor],
            batch_idx: int
    ):
        images, labels = batch
        reconstructions, _ = self.forward(images)
        reconstructions.squeeze_(1)
        self.fid(images, reconstructions)
        if batch_idx == 0:
            return {
                'real images': images,
                # 'real labels': labels,
            }
        return {
            'real images': None,
            # 'real labels': None
        }

    def validation_epoch_end(self, outputs: tp.List[tp.Any]) -> None:
        n_images = 5
        for output in outputs:
            if output['real images'] is not None:
                control_images = output['real images'][0:n_images]
                # control_labels = output['real labels'][0:n_images]
                break

        reconstructed, _ = self.forward(control_images)
        reconstructed.squeeze_(1)
        generated = self.generate_samples(n_images)
        self.log('fid score', self.fid)
        self.logger.experiment.log(
            {
                'generated images': wandb.Image(torchvision.utils.make_grid(generated, nrow=n_images)),
                'reconstructed images': wandb.Image(
                    torchvision.utils.make_grid(
                        torch.vstack((control_images, reconstructed)), nrow=n_images))
            },
        )

    def configure_optimizers(self):
        return instantiate(self.optimizer_conf, params=self.parameters())


def build_net(size, latent_dim=400):
    return nn.Sequential(nn.Linear(size, latent_dim), nn.ReLU(), nn.Linear(latent_dim, latent_dim), nn.ReLU(),
                         nn.Linear(latent_dim, size), nn.Tanh())


class RealNVP(LightningModule):
    def __init__(self, shape, maskn, latent_dim, optimizer_conf=None, *args, **kwargs):
        super(RealNVP, self).__init__(*args, **kwargs)

        # Create a flow
        # nets:  a function that returns a PyTorch neural network, e.g., nn.Sequential, s = nets(), s: dim(X) -> dim(X)
        # nett:  a function that returns a PyTorch neural network, e.g., nn.Sequential, t = nett(), t: dim(X) -> dim(X)
        # mask:  a torch.Tensor of size #number_of_coupling_layers x #dim(X)
        # prior: an object from torch.distributions e.g., torch.distributions.MultivariateNormal
        c, h, w = shape
        self.shape = tuple(shape)
        size = c * h * w
        self.prior_mu = nn.Parameter(torch.zeros(size), requires_grad=False)
        self.prior_sigma = nn.Parameter(torch.ones(size), requires_grad=False)
        m = torch.vstack(
            (
                torch.hstack(
                    (torch.zeros(size // 2), torch.ones(size // 2))
                ),
                torch.hstack(
                    (torch.ones(size // 2), torch.zeros(size // 2))
                )
            )
        )
        mask = m.repeat(maskn, 1)
        self.mask = nn.Parameter(mask, requires_grad=False)
        self.t = torch.nn.ModuleList([build_net(size, latent_dim) for _ in range(len(mask))])
        self.s = torch.nn.ModuleList([build_net(size, latent_dim) for _ in range(len(mask))])
        self.optimizer_conf = optimizer_conf
        self.fid = FidScore()

    def g1(self, y, mask, t, s):
        my = mask * y
        s_ = (1 - mask) * s(my)
        return my + (1 - mask) * (y - t(my)) * torch.exp(-s_), -s_.sum(dim=1)

    def g(self, z):
        log_det_J = 0
        x = z
        log_det_J += torch.sum(1 / x - 1 / (1 - x), dim=-1)
        x = torch.log(x / (1 - x))
        for mask, t, s in zip(reversed(self.mask), reversed(self.t), reversed(self.s)):
            x, lgdJ = self.g1(x, mask, t, s)
            log_det_J += lgdJ
        return x, log_det_J

    def f1(self, x, mask, t, s):
        mx = mask * x
        s_ = (1 - mask) * s(mx)
        return mx + (1 - mask) * (x * torch.exp(s_) + t(mx)), s_.sum(dim=1)

    def f(self, x):
        # Compute f(x) = z and log_det_Jakobian of f,
        #    where self.mask[i], self.t[i], self.s[i] define a i-th masked coupling layer
        # x: a torch.Tensor, of shape batchSize x dim(X), is a datapoint
        # return z: a torch.Tensor of shape batchSize x dim(X), a hidden representations
        # return log_det_J: a torch.Tensor of len batchSize
        z = x
        log_det_J = 0
        for mask, t, s in zip(self.mask, self.t, self.s):
            z, lgdJ = self.f1(z, mask, t, s)
            log_det_J += lgdJ
        log_det_J += torch.sum(z + 2 * torch.log1p(torch.exp(z)), dim=-1)
        z = F.sigmoid(z)
        return z, log_det_J

    def log_prob(self, x):
        # Compute and return log p(x)
        # using the change of variable formula and log_det_J computed by f
        # return logp: torch.Tensor of len batchSize
        z, log_det_J = self.f(x)
        lp = -0.5*(z**2).sum(-1)
        logp = lp + log_det_J
        return logp

    def training_step(self, batch, batch_idx, *args, **kwargs):
        batch, _ = batch
        batch = batch.flatten(1)
        ll = self.log_prob(batch).mean()
        self.log('log likelihood', ll)
        return -ll

    def forward(self, batch, *args, **kwargs):
        z, log_det_J = self.f(batch)
        return z

    @torch.no_grad()
    def generate_samples(self, num_samples=1):
        """
        Generate samples from the model.
        Input: num_samples, int - number of samples to generate.
        Return: Tensor of shape num_samples x D.
        """
        device = next(iter(self.parameters()))[0].device
        latent_samples = self.prior_mu + torch.randn((num_samples, *self.prior_mu.shape), device=device) * self.prior_sigma
        samples, _ = self.g(latent_samples)
        return samples

    def validation_step(
            self,
            batch: tp.Tuple[torch.Tensor, torch.Tensor],
            batch_idx: int
    ):
        images, labels = batch
        generated_samples = self.generate_samples(num_samples=images.shape[0])
        generated_samples = generated_samples.unflatten(dim=1,
                                                        sizes=self.shape).clamp(0, 1)
        self.fid(images, generated_samples)
        if batch_idx == 0:
            return {
                'real images': images,
                # 'real labels': labels,
            }
        return {
            'real images': None,
            # 'real labels': None
        }

    def validation_epoch_end(self, outputs: tp.List[tp.Any]) -> None:
        n_images = 5
        for output in outputs:
            if output['real images'] is not None:
                control_images = output['real images'][0:n_images]
                # control_labels = output['real labels'][0:n_images]
                break

        generated = self.generate_samples(n_images).unflatten(dim=1,
                                                              sizes=self.shape).clamp(0, 1)
        self.log('fid score', self.fid)
        # self.logger.experiment.log(
        #     {
        #         'generated images': wandb.Image(
        #             torchvision.utils.make_grid(
        #                 torch.vstack((control_images, generated)),
        #                 nrow=n_images)
        #         ),
        #     }
        # )

    def configure_optimizers(self):
        return instantiate(self.optimizer_conf, params=self.parameters())

import torch
from torch import nn
from torch.nn import functional as F
from math import log, pi, exp
import numpy as np
from scipy import linalg as la

logabs = lambda x: torch.log(torch.abs(x))


class ActNorm(nn.Module):
    def __init__(self, in_channel, logdet=True):
        super().__init__()

        self.loc = nn.Parameter(torch.zeros(1, in_channel, 1, 1))
        self.scale = nn.Parameter(torch.ones(1, in_channel, 1, 1))

        self.register_buffer("initialized", torch.tensor(0, dtype=torch.uint8))
        self.logdet = logdet

    def initialize(self, input):
        with torch.no_grad():
            flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
            mean = (
                flatten.mean(1)
                .unsqueeze(1)
                .unsqueeze(2)
                .unsqueeze(3)
                .permute(1, 0, 2, 3)
            )
            std = (
                flatten.std(1)
                .unsqueeze(1)
                .unsqueeze(2)
                .unsqueeze(3)
                .permute(1, 0, 2, 3)
            )

            self.loc.data.copy_(-mean)
            self.scale.data.copy_(1 / (std + 1e-6))

    def forward(self, input):
        _, _, height, width = input.shape

        if self.initialized.item() == 0:
            self.initialize(input)
            self.initialized.fill_(1)

        log_abs = logabs(self.scale)

        logdet = height * width * torch.sum(log_abs)

        if self.logdet:
            return self.scale * (input + self.loc), logdet

        else:
            return self.scale * (input + self.loc)

    def reverse(self, output):
        return output / self.scale - self.loc


class InvConv2d(nn.Module):
    def __init__(self, in_channel):
        super().__init__()

        weight = torch.randn(in_channel, in_channel)
        q, _ = torch.qr(weight)
        weight = q.unsqueeze(2).unsqueeze(3)
        self.weight = nn.Parameter(weight)

    def forward(self, input):
        _, _, height, width = input.shape

        out = F.conv2d(input, self.weight)
        logdet = (
            height * width * torch.slogdet(self.weight.squeeze().double())[1].float()
        )

        return out, logdet

    def reverse(self, output):
        return F.conv2d(
            output, self.weight.squeeze().inverse().unsqueeze(2).unsqueeze(3)
        )


class InvConv2dLU(nn.Module):
    def __init__(self, in_channel):
        super().__init__()

        weight = np.random.randn(in_channel, in_channel)
        q, _ = la.qr(weight)
        w_p, w_l, w_u = la.lu(q.astype(np.float32))
        w_s = np.diag(w_u)
        w_u = np.triu(w_u, 1)
        u_mask = np.triu(np.ones_like(w_u), 1)
        l_mask = u_mask.T

        w_p = torch.from_numpy(w_p)
        w_l = torch.from_numpy(w_l)
        w_s = torch.from_numpy(w_s)
        w_u = torch.from_numpy(w_u)

        self.register_buffer("w_p", w_p)
        self.register_buffer("u_mask", torch.from_numpy(u_mask))
        self.register_buffer("l_mask", torch.from_numpy(l_mask))
        self.register_buffer("s_sign", torch.sign(w_s))
        self.register_buffer("l_eye", torch.eye(l_mask.shape[0]))
        self.w_l = nn.Parameter(w_l)
        self.w_s = nn.Parameter(logabs(w_s))
        self.w_u = nn.Parameter(w_u)

    def forward(self, input):
        _, _, height, width = input.shape

        weight = self.calc_weight()

        out = F.conv2d(input, weight)
        logdet = height * width * torch.sum(self.w_s)

        return out, logdet

    def calc_weight(self):
        weight = (
            self.w_p
            @ (self.w_l * self.l_mask + self.l_eye)
            @ ((self.w_u * self.u_mask) + torch.diag(self.s_sign * torch.exp(self.w_s)))
        )

        return weight.unsqueeze(2).unsqueeze(3)

    def reverse(self, output):
        weight = self.calc_weight()

        return F.conv2d(output, weight.squeeze().inverse().unsqueeze(2).unsqueeze(3))


class ZeroConv2d(nn.Module):
    def __init__(self, in_channel, out_channel, padding=1):
        super().__init__()

        self.conv = nn.Conv2d(in_channel, out_channel, 3, padding=0)
        self.conv.weight.data.zero_()
        self.conv.bias.data.zero_()
        self.scale = nn.Parameter(torch.zeros(1, out_channel, 1, 1))

    def forward(self, input):
        out = F.pad(input, [1, 1, 1, 1], value=1)
        out = self.conv(out)
        out = out * torch.exp(self.scale * 3)

        return out


class AffineCoupling(nn.Module):
    def __init__(self, in_channel, filter_size=512, affine=True):
        super().__init__()

        self.affine = affine

        self.net = nn.Sequential(
            nn.Conv2d(in_channel // 2, filter_size, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(filter_size, filter_size, 1),
            nn.ReLU(inplace=True),
            ZeroConv2d(filter_size, in_channel if self.affine else in_channel // 2),
        )

        self.net[0].weight.data.normal_(0, 0.05)
        self.net[0].bias.data.zero_()

        self.net[2].weight.data.normal_(0, 0.05)
        self.net[2].bias.data.zero_()

    def forward(self, input):
        in_a, in_b = input.chunk(2, 1)

        if self.affine:
            log_s, t = self.net(in_a).chunk(2, 1)
            # s = torch.exp(log_s)
            s = F.sigmoid(log_s + 2)
            # out_a = s * in_a + t
            out_b = (in_b + t) * s

            logdet = torch.sum(torch.log(s).view(input.shape[0], -1), 1)

        else:
            net_out = self.net(in_a)
            out_b = in_b + net_out
            logdet = None

        return torch.cat([in_a, out_b], 1), logdet

    def reverse(self, output):
        out_a, out_b = output.chunk(2, 1)

        if self.affine:
            log_s, t = self.net(out_a).chunk(2, 1)
            # s = torch.exp(log_s)
            s = F.sigmoid(log_s + 2)
            # in_a = (out_a - t) / s
            in_b = out_b / s - t

        else:
            net_out = self.net(out_a)
            in_b = out_b - net_out

        return torch.cat([out_a, in_b], 1)


class Flow(nn.Module):
    def __init__(self, in_channel, affine=True, conv_lu=True):
        super().__init__()

        self.actnorm = ActNorm(in_channel)

        if conv_lu:
            self.invconv = InvConv2dLU(in_channel)

        else:
            self.invconv = InvConv2d(in_channel)

        self.coupling = AffineCoupling(in_channel, affine=affine)

    def forward(self, input):
        out, logdet = self.actnorm(input)
        out, det1 = self.invconv(out)
        out, det2 = self.coupling(out)

        logdet = logdet + det1
        if det2 is not None:
            logdet = logdet + det2

        return out, logdet

    def reverse(self, output):
        input = self.coupling.reverse(output)
        input = self.invconv.reverse(input)
        input = self.actnorm.reverse(input)

        return input


def gaussian_log_p(x, mean, log_sd):
    return -0.5 * log(2 * pi) - log_sd - 0.5 * (x - mean) ** 2 / torch.exp(2 * log_sd)


def gaussian_sample(eps, mean, log_sd):
    return mean + torch.exp(log_sd) * eps


class Block(nn.Module):
    def __init__(self, in_channel, n_flow, split=True, affine=True, conv_lu=True):
        super().__init__()

        squeeze_dim = in_channel * 4

        self.flows = nn.ModuleList()
        for i in range(n_flow):
            self.flows.append(Flow(squeeze_dim, affine=affine, conv_lu=conv_lu))

        self.split = split

        if split:
            self.prior = ZeroConv2d(in_channel * 2, in_channel * 4)

        else:
            self.prior = ZeroConv2d(in_channel * 4, in_channel * 8)

    def forward(self, input):
        b_size, n_channel, height, width = input.shape
        squeezed = input.view(b_size, n_channel, height // 2, 2, width // 2, 2)
        squeezed = squeezed.permute(0, 1, 3, 5, 2, 4)
        out = squeezed.contiguous().view(b_size, n_channel * 4, height // 2, width // 2)

        logdet = 0

        for flow in self.flows:
            out, det = flow(out)
            logdet = logdet + det

        if self.split:
            out, z_new = out.chunk(2, 1)
            mean, log_sd = self.prior(out).chunk(2, 1)
            log_p = gaussian_log_p(z_new, mean, log_sd)
            log_p = log_p.view(b_size, -1).sum(1)

        else:
            zero = torch.zeros_like(out)
            mean, log_sd = self.prior(zero).chunk(2, 1)
            log_p = gaussian_log_p(out, mean, log_sd)
            log_p = log_p.view(b_size, -1).sum(1)
            z_new = out

        return out, logdet, log_p, z_new

    def reverse(self, output, eps=None, reconstruct=False):
        input = output

        if reconstruct:
            if self.split:
                input = torch.cat([output, eps], 1)

            else:
                input = eps

        else:
            if self.split:
                mean, log_sd = self.prior(input).chunk(2, 1)
                z = gaussian_sample(eps, mean, log_sd)
                input = torch.cat([output, z], 1)

            else:
                zero = torch.zeros_like(input)
                # zero = F.pad(zero, [1, 1, 1, 1], value=1)
                mean, log_sd = self.prior(zero).chunk(2, 1)
                z = gaussian_sample(eps, mean, log_sd)
                input = z

        for flow in self.flows[::-1]:
            input = flow.reverse(input)

        b_size, n_channel, height, width = input.shape

        unsqueezed = input.view(b_size, n_channel // 4, 2, 2, height, width)
        unsqueezed = unsqueezed.permute(0, 1, 4, 2, 5, 3)
        unsqueezed = unsqueezed.contiguous().view(
            b_size, n_channel // 4, height * 2, width * 2
        )

        return unsqueezed


class Glow(nn.Module):
    def __init__(
        self, in_channel, n_flow, n_block, affine=True, conv_lu=True
    ):
        super().__init__()

        self.blocks = nn.ModuleList()
        n_channel = in_channel
        for i in range(n_block - 1):
            self.blocks.append(Block(n_channel, n_flow, affine=affine, conv_lu=conv_lu))
            n_channel *= 2
        self.blocks.append(Block(n_channel, n_flow, split=False, affine=affine))

    def forward(self, input):
        log_p_sum = 0
        logdet = 0
        out = input
        z_outs = []

        for block in self.blocks:
            out, det, log_p, z_new = block(out)
            z_outs.append(z_new)
            logdet = logdet + det

            if log_p is not None:
                log_p_sum = log_p_sum + log_p

        return log_p_sum, logdet, z_outs

    def reverse(self, z_list, reconstruct=False):
        for i, block in enumerate(self.blocks[::-1]):
            if i == 0:
                input = block.reverse(z_list[-1], z_list[-1], reconstruct=reconstruct)

            else:
                input = block.reverse(input, z_list[-(i + 1)], reconstruct=reconstruct)

        return input