import torch.nn as nn

__all__ = ['SqueezeLayer', 'squeeze', 'unsqueeze']


class SqueezeLayer(nn.Module):
    def __init__(self, downscale_factor):
        super(SqueezeLayer, self).__init__()
        self.downscale_factor = downscale_factor

    def forward(self, x, logpx=None, reg_states=tuple(), reverse=False):
        if reverse:
            return self._upsample(x, logpx, reg_states)
        else:
            return self._downsample(x, logpx, reg_states)

    def _downsample(self, x, logpx=None, reg_states=tuple()):
        squeeze_x = squeeze(x, self.downscale_factor)
        return squeeze_x, logpx, reg_states

    def _upsample(self, y, logpy=None, reg_states=tuple()):
        unsqueeze_y = unsqueeze(y, self.downscale_factor)
        return unsqueeze_y, logpy, reg_states


def unsqueeze(input, upscale_factor=2):
    '''
    [:, C*r^2, H, W] -> [:, C, H*r, W*r]
    '''
    batch_size, in_channels, in_height, in_width = input.size()
    out_channels = in_channels // (upscale_factor**2)

    out_height = in_height * upscale_factor
    out_width = in_width * upscale_factor

    input_view = input.contiguous().view(batch_size, out_channels, upscale_factor, upscale_factor, in_height, in_width)

    output = input_view.permute(0, 1, 4, 2, 5, 3).contiguous()
    return output.view(batch_size, out_channels, out_height, out_width)


def squeeze(input, downscale_factor=2):
    '''
    [:, C, H*r, W*r] -> [:, C*r^2, H, W]
    '''
    batch_size, in_channels, in_height, in_width = input.size()
    out_channels = in_channels * (downscale_factor**2)

    out_height = in_height // downscale_factor
    out_width = in_width // downscale_factor

    input_view = input.contiguous().view(
        batch_size, in_channels, out_height, downscale_factor, out_width, downscale_factor
    )

    output = input_view.permute(0, 1, 3, 5, 2, 4).contiguous()
    return output.view(batch_size, out_channels, out_height, out_width)
