# -*- coding: UTF-8 -*-
import matplotlib.pyplot as plt
import torch
import torch.nn as nn


class Net(nn.Module):
    def __init__(self, upscale_factor):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 64, (5, 5), (1, 1), (2, 2))
        self.conv2 = nn.Conv2d(64, 32, (3, 3), (1, 1), (1, 1))
        self.conv3 = nn.Conv2d(32, 1 * (upscale_factor ** 2), (3, 3), (1, 1), (1, 1))
        self.pixel_shuffle = nn.PixelShuffle(upscale_factor)

    def forward(self, x):
        x = torch.tanh(self.conv1(x))
        x = torch.tanh(self.conv2(x))
        x = torch.sigmoid(self.pixel_shuffle(self.conv3(x)))
        return x


if __name__ == "__main__":
    model = Net(upscale_factor=3)
    print(model)
    # 用此处代码测试时要改模型，网络的input_channel = 3
    oritensor = torch.randn(1, 3, 33, 33)
    oritensor = torch.clamp(oritensor, 0, 1)
    newtensor = torch.clamp(model(oritensor), 0, 1)
    orinumpy = oritensor.detach().squeeze().permute(1, 2, 0).numpy()
    newnumpy = newtensor.detach().squeeze().permute(1, 2, 0).numpy()
    plt.imshow(orinumpy)
    plt.title('original noise')
    plt.show()
    plt.imshow(newnumpy)
    plt.title('reconstruct noise')
    plt.show()
    print(orinumpy.shape)
    print(newnumpy.shape)