from matplotlib import pyplot as plt
import torch
import torchvision as tv

samples0, samples1 = 60000, 10000

source0 = tv.datasets.MNIST("../MNIST", train = True, download = True)
source1 = tv.datasets.MNIST("../MNIST", train = False, download = True)
DATA0 = source0.data.flatten(1).float().cuda() / 255.
DATA1 = source1.data.flatten(1).float().cuda() / 255.
TARGET0 = DATA0
TARGET1 = DATA1

model = torch.nn.Sequential(
    torch.nn.Linear(28 * 28, 16 * 16),
    torch.nn.Tanh(),
    torch.nn.Linear(16 * 16, 8 * 8),
    torch.nn.Tanh(),
    torch.nn.Linear(8 * 8, 4 * 4),
    torch.nn.Tanh(),
    torch.nn.Linear(4 * 4, 8 * 8),
    torch.nn.Tanh(),
    torch.nn.Linear(8 * 8, 16 * 16),
    torch.nn.Tanh(),
    torch.nn.Linear(16 * 16, 28 * 28),
    torch.nn.Sigmoid()).cuda()
model.load_state_dict(torch.load("anomaly.pt"))

DATA = DATA1[:10]
TARGET = TARGET1[:10]
VALUE = model(DATA).detach()
LOSS = torch.nn.functional.mse_loss(VALUE, TARGET, reduction = "none").mean(1)
print(LOSS, flush = True)

DATA = DATA.reshape(10, 1, 28, 28).cpu()
VALUE = VALUE.reshape(10, 1, 28, 28).cpu()
GRID = tv.utils.make_grid(torch.cat([DATA, VALUE]), nrow = 10)
plt.imshow(GRID.permute(1, 2, 0))

plt.show()
