import numpy as np
import os
import torch
import torch.optim as optim
from torch.nn import functional as F
from torch.utils.data import DataLoader, TensorDataset, random_split
from torchvision.utils import save_image
from model.Multitask_AE import MAE, MGVAE, MLVAE, MGMVAE, MLMVAE
import sys
sys.path.append('..')
import random
import scipy.io as sio
from torchinfo import summary
from sklearn.metrics import roc_curve
import cv2


# runner for Multi-task AE
class MAE_runner():
    def __init__(self, args):
        self.args = args
        cuda_flag = not self.args.no_cuda and torch.cuda.is_available()
        self.device = torch.device("cuda" if cuda_flag else "cpu")
        torch.manual_seed(self.args.seed)
        random.seed(self.args.seed)
        self.path = self.args.run + 'MAE/'
        if not os.path.exists(self.path):
            os.makedirs(self.path)

    def get_optimizer(self, parameters):
        if self.args.optimizer == 'Adam':
            return optim.Adam(parameters, lr=self.args.lr, weight_decay=self.args.weight_decay, betas=(0.9, 0.999), eps=1e-7)
        elif self.args.optimizer == 'RMSProp':
            return optim.RMSprop(parameters, lr=self.args.lr, weight_decay=self.args.weight_decay)
        elif self.args.optimizer == 'SGD':
            return optim.SGD(parameters, lr=self.args.lr, momentum=0.9)
        else:
            raise NotImplementedError('Optimizer {} not understood.'.format(self.args.optimizer))

    def load_data(self, noise_level='all'):
        dataset, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min = self.load_training_data()
        train_size = int(len(dataset)*0.9)
        val_size = int(len(dataset)*0.1)
        dataset_train, dataset_val = random_split(dataset, [train_size, val_size],
                                                  generator=torch.Generator().manual_seed(self.args.seed))
        dataset_test = self.load_testing_data(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        dataset_test1 = self.load_testing_img1(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        dataset_test2 = self.load_testing_img2(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        train_loader = DataLoader(
            dataset_train,
            batch_size=self.args.batch_size, shuffle=True, num_workers=2)
        val_loader = DataLoader(
            dataset_val,
            batch_size=self.args.batch_size, shuffle=True, num_workers=2)
        test_loader = DataLoader(
            dataset_test,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        test_loader1 = DataLoader(
            dataset_test1,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        test_loader2 = DataLoader(
            dataset_test2,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        print("Data Loaded!")
        return train_loader, val_loader, test_loader, test_loader1, test_loader2

    def load_training_data(self):
        # Load training data
        training_data = sio.loadmat('./data/training_data.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        train_input_RSS = training_data['RSS']
        RSS_max = np.amax(train_input_RSS, axis=0, keepdims=True)
        RSS_min = np.amin(train_input_RSS, axis=0, keepdims=True)
        train_input_RSS = (train_input_RSS - RSS_min) / (RSS_max - RSS_min)
        train_input_RSS = train_input_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        train_denoise_RSS = training_data['RSS_denoise']
        RSS_denoise_max = np.amax(train_denoise_RSS, axis=0, keepdims=True)
        RSS_denoise_min = np.amin(train_denoise_RSS, axis=0, keepdims=True)
        train_denoise_RSS = (train_denoise_RSS - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        train_denoise_RSS = train_denoise_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        train_slf_img = training_data['slf_img']
        train_slf_img = train_slf_img.reshape(-1, self.args.K0*self.args.K1)
        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        train_ab = training_data['ab_norm']
        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        train_noise_level = training_data['sig_epsilon_class']
        train_noise_level = np.squeeze(train_noise_level)

        train_input_RSS = torch.from_numpy(train_input_RSS).float()
        train_denoise_RSS = torch.from_numpy(train_denoise_RSS).float()
        train_slf_img = torch.from_numpy(train_slf_img).float()
        train_ab = torch.from_numpy(train_ab).float()
        train_noise_level = torch.from_numpy(train_noise_level).long()

        Train_Dataset = TensorDataset(train_input_RSS, train_denoise_RSS, train_slf_img, train_ab, train_noise_level)
        return Train_Dataset, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min

    def load_testing_data(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data
        testing_data = sio.loadmat('./data/testing_data.mat')
        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_class = testing_data['sig_epsilon_class']
        test_noise_class = np.squeeze(test_noise_class)
        idx = (test_noise_class == 0) | (test_noise_class == 1) | (test_noise_class == 2)
        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            idx = (test_noise_class == 0)
        elif noise_level == 'mid':
            idx = (test_noise_class == 1)
        elif noise_level == 'high':
            idx = (test_noise_class == 2)
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))
        test_noise_class = test_noise_class[idx]
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS = testing_data['RSS'][idx]
        test_input_RSS = (test_input_RSS - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS = test_input_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS = testing_data['RSS_denoise'][idx]
        test_denoise_RSS = (test_denoise_RSS - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS = test_denoise_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img = testing_data['slf_img'][idx]
        test_slf_img = test_slf_img.reshape(-1, self.args.K0*self.args.K1)

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab = testing_data['ab_norm'][idx]

        test_input_RSS = torch.from_numpy(test_input_RSS).float()
        test_denoise_RSS = torch.from_numpy(test_denoise_RSS).float()
        test_slf_img = torch.from_numpy(test_slf_img).float()
        test_ab = torch.from_numpy(test_ab).float()
        test_noise_class = torch.from_numpy(test_noise_class).long()

        Test_Dataset = TensorDataset(test_input_RSS, test_denoise_RSS, test_slf_img, test_ab, test_noise_class)
        return Test_Dataset

    def load_testing_img1(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data for image 1
        testing_data1 = sio.loadmat('./data/testing_data_image1.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS1 = testing_data1['RSS']
        test_input_RSS1 = (test_input_RSS1 - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS1 = test_input_RSS1.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS1 = testing_data1['RSS_denoise']
        test_denoise_RSS1 = (test_denoise_RSS1 - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS1 = test_denoise_RSS1.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img1 = testing_data1['slf_img']
        test_slf_img1 = test_slf_img1.reshape(1, self.args.K0*self.args.K1)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        test_slf_img1 = np.tile(test_slf_img1, (test_input_RSS1.shape[0], 1))

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab1 = testing_data1['ab_norm']

        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_level1 = testing_data1['sig_epsilon_class']
        test_noise_level1 = np.squeeze(test_noise_level1)

        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            test_input_RSS1 = test_input_RSS1[:500]
            test_denoise_RSS1 = test_denoise_RSS1[:500]
            test_slf_img1 = test_slf_img1[:500]
            test_ab1 = test_ab1[:500]
            test_noise_level1 = test_noise_level1[:500]
        elif noise_level == 'mid':
            test_input_RSS1 = test_input_RSS1[500:1000]
            test_denoise_RSS1 = test_denoise_RSS1[500:1000]
            test_slf_img1 = test_slf_img1[500:1000]
            test_ab1 = test_ab1[500:1000]
            test_noise_level1 = test_noise_level1[500:1000]
        elif noise_level == 'high':
            test_input_RSS1 = test_input_RSS1[1000:]
            test_denoise_RSS1 = test_denoise_RSS1[1000:]
            test_slf_img1 = test_slf_img1[1000:]
            test_ab1 = test_ab1[1000:]
            test_noise_level1 = test_noise_level1[1000:]
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))

        test_input_RSS1 = torch.from_numpy(test_input_RSS1).float()
        test_denoise_RSS1 = torch.from_numpy(test_denoise_RSS1).float()
        test_slf_img1 = torch.from_numpy(test_slf_img1).float()
        test_ab1 = torch.from_numpy(test_ab1).float()
        test_noise_level1 = torch.from_numpy(test_noise_level1).long()

        Test_Dataset_1 = TensorDataset(test_input_RSS1, test_denoise_RSS1, test_slf_img1, test_ab1, test_noise_level1)
        return Test_Dataset_1

    def load_testing_img2(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data for image 2
        testing_data2 = sio.loadmat('./data/testing_data_image2.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS2 = testing_data2['RSS']
        test_input_RSS2 = (test_input_RSS2 - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS2 = test_input_RSS2.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS2 = testing_data2['RSS_denoise']
        test_denoise_RSS2 = (test_denoise_RSS2 - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS2 = test_denoise_RSS2.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img2 = testing_data2['slf_img']
        test_slf_img2 = test_slf_img2.reshape(1, self.args.K0 * self.args.K1)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        test_slf_img2 = np.tile(test_slf_img2, (test_input_RSS2.shape[0], 1))

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab2 = testing_data2['ab_norm']

        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_level2 = testing_data2['sig_epsilon_class']
        test_noise_level2 = np.squeeze(test_noise_level2)

        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            test_input_RSS2 = test_input_RSS2[:500]
            test_denoise_RSS2 = test_denoise_RSS2[:500]
            test_slf_img2 = test_slf_img2[:500]
            test_ab2 = test_ab2[:500]
            test_noise_level2 = test_noise_level2[:500]
        elif noise_level == 'mid':
            test_input_RSS2 = test_input_RSS2[500:1000]
            test_denoise_RSS2 = test_denoise_RSS2[500:1000]
            test_slf_img2 = test_slf_img2[500:1000]
            test_ab2 = test_ab2[500:1000]
            test_noise_level2 = test_noise_level2[500:1000]
        elif noise_level == 'high':
            test_input_RSS2 = test_input_RSS2[1000:]
            test_denoise_RSS2 = test_denoise_RSS2[1000:]
            test_slf_img2 = test_slf_img2[1000:]
            test_ab2 = test_ab2[1000:]
            test_noise_level2 = test_noise_level2[1000:]
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))

        test_input_RSS2 = torch.from_numpy(test_input_RSS2).float()
        test_denoise_RSS2 = torch.from_numpy(test_denoise_RSS2).float()
        test_slf_img2 = torch.from_numpy(test_slf_img2).float()
        test_ab2 = torch.from_numpy(test_ab2).float()
        test_noise_level2 = torch.from_numpy(test_noise_level2).long()

        Test_Dataset_2 = TensorDataset(test_input_RSS2, test_denoise_RSS2, test_slf_img2, test_ab2, test_noise_level2)
        return Test_Dataset_2

    def train(self, model, train_loader, optimizer, epoch):
        model.train()
        train_loss = 0                  # total loss
        train_loss_task0 = 0            # loss for reconstruct RSS
        train_loss_task1 = 0            # loss for SLF image estimation
        train_loss_task2 = 0            # loss for parameters estimation
        train_loss_task3 = 0            # loss for noise level prediction
        train_acc = 0                   # accuracy for noise level prediction
        for batch_idx,  (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(train_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            optimizer.zero_grad()
            RSS, SLF, ab, noise_level = model(input_RSS)
            loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum')/(self.args.N * self.args.P * self.args.P)
            loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum')/(self.args.K0 * self.args.K1)
            loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum')/(self.args.N + 1)
            loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
            loss = self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1 \
                   + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3
            train_loss += loss.item()
            train_loss_task0 += loss_task0.item()
            train_loss_task1 += loss_task1.item()
            train_loss_task2 += loss_task2.item()
            train_loss_task3 += loss_task3.item()
            _, noise_pred = torch.max(noise_level, 1)
            acc_task3 = (noise_pred == target_noise_level).sum()
            train_acc += acc_task3.item()
            loss /= input_RSS.size(0)
            loss.backward()
            optimizer.step()

            if batch_idx % self.args.log_interval == 0:
                print(
                    'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tLoss0: {:.6f}\tLoss1: {:.6f}\tLoss2: {:.6f}\tLoss3: {:.6f}\tAcc: {:.6f}'.format(
                        epoch, batch_idx * len(input_RSS), len(train_loader.dataset),
                               100. * batch_idx / len(train_loader),
                               loss.item(),
                               loss_task0.item() / len(input_RSS), loss_task1.item() / len(input_RSS),
                               loss_task2.item() / len(input_RSS), loss_task3.item() / len(input_RSS), acc_task3.item() / len(input_RSS)))

        train_loss /= len(train_loader.dataset)
        train_loss_task0 /= len(train_loader.dataset)
        train_loss_task1 /= len(train_loader.dataset)
        train_loss_task2 /= len(train_loader.dataset)
        train_loss_task3 /= len(train_loader.dataset)
        train_acc /= len(train_loader.dataset)
        print('====> Epoch: {} Average loss: {:.6f}'.format(epoch, train_loss))
        print('====> RSS loss: {:.6f}'.format(train_loss_task0))
        print('====> SLF loss: {:.6f}'.format(train_loss_task1))
        print('====> ab loss: {:.6f}'.format(train_loss_task2))
        print('====> noise level loss: {:.6f}'.format(train_loss_task3))
        print('====> noise level Accuracy: {:.4f}'.format(train_acc))
        return train_loss, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3, train_acc

    def validate(self, model, val_loader):
        model.eval()
        val_loss = 0                  # total loss
        val_rmse_slf = 0              # slf reconstruction loss
        val_mae_slf = 0              # slf reconstruction loss
        for batch_idx,  (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(val_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            RSS, SLF, ab, noise_level = model(input_RSS)
            loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum')/(self.args.N * self.args.P * self.args.P)
            loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum')/(self.args.K0 * self.args.K1)
            loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum')/(self.args.N + 1)
            loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
            loss = self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1 \
                   + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3
            val_loss += loss.item()
            slf_mse = F.mse_loss(SLF, target_SLF, reduction='sum')/(self.args.K0 * self.args.K1)
            slf_mae = F.l1_loss(SLF, target_SLF, reduction='sum')/(self.args.K0 * self.args.K1)
            val_rmse_slf += slf_mse.item()
            val_mae_slf += slf_mae.item()
        val_loss /= len(val_loader.dataset)
        val_rmse_slf /= len(val_loader.dataset)
        val_rmse_slf = np.sqrt(val_rmse_slf)
        val_mae_slf /= len(val_loader.dataset)
        print('====> Validation set loss: {:.6f}'.format(val_loss))
        print('====> Validation set SLF rmse: {:.6f}'.format(val_rmse_slf))
        print('====> Validation set SLF mae: {:.6f}'.format(val_mae_slf))
        return val_loss, val_rmse_slf, val_mae_slf

    def test(self, model, test_loader):
        model.eval()
        test_loss = 0  # total loss
        test_loss_task0 = 0  # loss for reconstruct RSS
        test_loss_task1 = 0  # loss for SLF image estimation
        test_loss_task2 = 0  # loss for parameters estimation
        test_loss_task3 = 0  # loss for noise level prediction
        with torch.no_grad():
            for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader):
                input_RSS = input_RSS.to(self.device)
                target_RSS = target_RSS.to(self.device)
                target_SLF = target_SLF.to(self.device)
                target_ab = target_ab.to(self.device)
                target_noise_level = target_noise_level.to(self.device)
                RSS, SLF, ab, noise_level = model(input_RSS)
                loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum') / (
                            self.args.N * self.args.P * self.args.P)
                loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
                loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum') / (self.args.N + 1)
                loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
                loss = self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1 \
                       + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3

                test_loss += loss.item()
                test_loss_task0 += loss_task0.item()
                test_loss_task1 += loss_task1.item()
                test_loss_task2 += loss_task2.item()
                test_loss_task3 += loss_task3.item()

        test_loss /= len(test_loader.dataset)
        test_loss_task0 /= len(test_loader.dataset)
        test_loss_task1 /= len(test_loader.dataset)
        test_loss_task2 /= len(test_loader.dataset)
        test_loss_task3 /= len(test_loader.dataset)
        test_acc = self.test_accuracy(model, test_loader)
        print('====> Test set loss: {:.6f}'.format(test_loss))
        print('====> Test set RSS loss: {:.6f}'.format(test_loss_task0))
        print('====> Test set SLF loss: {:.6f}'.format(test_loss_task1))
        print('====> Test set ab loss: {:.6f}'.format(test_loss_task2))
        print('====> Test set noise level loss: {:.6f}'.format(test_loss_task3))
        print('====> Test set noise level Accuracy: {:.4f}'.format(test_acc))
        return test_loss, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc

    def train_save(self):
        model = MAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))

        def weights_init(m):
            if isinstance(m, torch.nn.Linear) or isinstance(m, torch.nn.Conv2d) or isinstance(m, torch.nn.ConvTranspose2d):
                torch.nn.init.xavier_uniform_(m.weight.data)
                if m.bias is not None:
                    torch.nn.init.zeros_(m.bias.data)
        model.apply(weights_init)
        model = model.to(self.device)
        summary(model, input_size=(self.args.batch_size, self.args.N, self.args.P, self.args.P))
        train_loader, val_loader, test_loader, test_loader1, test_loader2 = self.load_data()
        optimizer = self.get_optimizer(model.parameters())
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.25, patience=10, min_lr=1e-5)
        train_set = []
        val_set = []
        test_set = []       # for test set
        test_set1 = []      # for test image 1
        test_set2 = []      # for test image 2
        train_loss_path = self.path + 'train_loss_' + str(self.args.n_epochs) + '.npy'
        val_loss_path = self.path + 'val_loss_' + str(self.args.n_epochs) + '.npy'
        test_loss_path = self.path + 'test_loss_' + str(self.args.n_epochs) + '.npz'
        model_path = self.path + 'model_' + str(self.args.n_epochs) + '.pth'
        current_val_slf_loss = np.inf
        for epoch in range(1, self.args.n_epochs+1):
            print("Epoch %d learning rate：%f" % (epoch, optimizer.param_groups[0]['lr']))
            train_loss, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3, train_acc = self.train(model,
                                                                                                            train_loader,
                                                                                                            optimizer,
                                                                                                            epoch)
            train_set.append([train_loss, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3, train_acc])
            print('====> Validation Loss for Epoch {:d}'.format(epoch))
            val_loss, val_slf_rmse, val_slf_mae = self.validate(model, val_loader)
            val_set.append([val_loss, val_slf_rmse, val_slf_mae])
            scheduler.step(val_loss)
            print('====> Test for Epoch {:d}'.format(epoch))
            test_loss, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc = self.test(model,
                                                                                                                test_loader)
            test_set.append([test_loss, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc])
            print('====> Test Image1 for Epoch {:d}'.format(epoch))
            test_loss, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc = self.test(model,
                                                                                                                test_loader1)
            test_set1.append([test_loss, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc])
            print('====> Test Image2 for Epoch {:d}'.format(epoch))
            test_loss, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc = self.test(model,
                                                                                                                test_loader2)
            test_set2.append([test_loss, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc])
            if epoch % self.args.save_freq == 0:
                ckpt_path = self.path + '/model_' + str(epoch) + '.pth'
                torch.save(model.state_dict(), ckpt_path)
                print('checkpoint{}.pth saved!'.format(epoch))

            if val_slf_mae < current_val_slf_loss:
                torch.save(model.state_dict(), model_path)
                current_val_slf_loss = val_slf_mae

        train_loss = np.asarray(train_set).reshape(-1, 6)
        val_loss = np.asarray(val_set).reshape(-1, 3)
        test_loss = np.asarray(test_set).reshape(-1, 6)
        test_loss1 = np.asarray(test_set1).reshape(-1, 6)
        test_loss2 = np.asarray(test_set2).reshape(-1, 6)
        np.save(train_loss_path, train_loss)
        np.save(val_loss_path, val_loss)
        np.savez(test_loss_path, test_loss=test_loss, test_loss1=test_loss1, test_loss2=test_loss2)

    def test_accuracy(self, model, test_loader):
        model.eval()
        acc = 0
        for i, (input_RSS, *_, target_noise_level) in enumerate(test_loader):
            input_RSS = input_RSS.to(self.device)
            *_, noise_pred = model(input_RSS)  # noise_pred: [batch_size, 3]
            noise_pred = noise_pred.detach().cpu().numpy()
            noise_pred = np.argmax(noise_pred, axis=1)
            labels = target_noise_level.detach().cpu().numpy()
            acc += np.sum((noise_pred == labels).astype(int))
        acc /= len(test_loader.dataset)
        return acc

    def test_model(self, noise_level='all'):
        model = MAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))
        model = model.to(self.device)
        model_path = self.path + 'model_' + str(self.args.n_epochs) + '.pth'
        model.load_state_dict(torch.load(model_path, map_location=self.device))
        model.eval()
        *_, test_loader, test_loader1, test_loader2 = self.load_data(noise_level)
        noise_class = 0  # noise_class = 0 stands for all noise classes
        if noise_level == 'low':
            noise_class = 1
        elif noise_level == 'mid':
            noise_class = 2
        elif noise_level == 'high':
            noise_class = 3

        # test model for testing set
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(
                test_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            RSS, SLF, ab, noise_pred = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader.dataset)
        ab_rmse /= len(test_loader.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader)
        print('====> Test for Testing Set:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))

        # # test model for image 1
        # slf_rmse = 0
        # slf_mae = 0
        # ab_rmse = 0
        # for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader1):
        #     input_RSS = input_RSS.to(self.device)
        #     target_RSS = target_RSS.to(self.device)
        #     target_SLF = target_SLF.to(self.device)
        #     target_ab = target_ab.to(self.device)
        #     target_noise_level = target_noise_level.to(self.device)
        #     RSS, SLF, ab, noise_pred = model(input_RSS)
        #     # plot SLF reconstruction image for image1
        #     n = min(target_RSS.size(0), 8)
        #     comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
        #                             SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
        #     save_image(comparison.cpu(),
        #                self.path + 'reconstruction_SLF1_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
        #                nrow=n)
        #     # test SLF image RMSE
        #     slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
        #     # test SLF image MAE
        #     slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
        #     # test ab parameters RMSE
        #     ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)
        #
        # slf_rmse /= len(test_loader1.dataset)
        # slf_rmse = np.sqrt(slf_rmse)
        # slf_mae /= len(test_loader1.dataset)
        # ab_rmse /= len(test_loader1.dataset)
        # ab_rmse = np.sqrt(ab_rmse)
        # noise_acc = self.test_accuracy(model, test_loader1)
        # print('====> Test for image 1:')
        # print('====> Noise_level: ' + noise_level)
        # print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        # print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        # print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        # print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))
        #
        # # test model for image 2
        # slf_rmse = 0
        # slf_mae = 0
        # ab_rmse = 0
        # for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader2):
        #     input_RSS = input_RSS.to(self.device)
        #     target_RSS = target_RSS.to(self.device)
        #     target_SLF = target_SLF.to(self.device)
        #     target_ab = target_ab.to(self.device)
        #     target_noise_level = target_noise_level.to(self.device)
        #     RSS, SLF, ab, noise_pred = model(input_RSS)
        #     # plot SLF reconstruction image for image1
        #     n = min(target_RSS.size(0), 8)
        #     comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
        #                             SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
        #     save_image(comparison.cpu(),
        #                self.path + 'reconstruction_SLF2_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
        #                nrow=n)
        #     # test SLF image RMSE
        #     slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
        #     # test SLF image MAE
        #     slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
        #     # test ab parameters RMSE
        #     ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)
        #
        # slf_rmse /= len(test_loader2.dataset)
        # slf_rmse = np.sqrt(slf_rmse)
        # slf_mae /= len(test_loader2.dataset)
        # ab_rmse /= len(test_loader2.dataset)
        # ab_rmse = np.sqrt(ab_rmse)
        # noise_acc = self.test_accuracy(model, test_loader2)
        # print('====> Test for image 2:')
        # print('====> Noise_level: ' + noise_level)
        # print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        # print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        # print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        # print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))

    def load_data_metrics(self, noise_level='all'):
        # Load training data
        training_data = sio.loadmat('./data/training_data.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        train_input_RSS = training_data['RSS']
        RSS_max = np.amax(train_input_RSS, axis=0, keepdims=True)
        RSS_min = np.amin(train_input_RSS, axis=0, keepdims=True)
        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        train_denoise_RSS = training_data['RSS_denoise']
        RSS_denoise_max = np.amax(train_denoise_RSS, axis=0, keepdims=True)
        RSS_denoise_min = np.amin(train_denoise_RSS, axis=0, keepdims=True)

        # Load testing data
        testing_data = sio.loadmat('./data/testing_samples1.mat') # sio.loadmat('./data/testing_data_YOLO/mat_file2/testing_YOLO_data.mat')
        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_class = testing_data['sig_epsilon_class']
        test_noise_class = np.squeeze(test_noise_class)
        idx = (test_noise_class == 0) | (test_noise_class == 1) | (test_noise_class == 2)
        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            idx = (test_noise_class == 0)
        elif noise_level == 'mid':
            idx = (test_noise_class == 1)
        elif noise_level == 'high':
            idx = (test_noise_class == 2)
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))
        test_noise_class = test_noise_class[idx]
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS = testing_data['RSS'][idx]
        test_input_RSS = (test_input_RSS - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS = test_input_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS = testing_data['RSS_denoise'][idx]
        test_denoise_RSS = (test_denoise_RSS - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS = test_denoise_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img = testing_data['slf_img'][idx]
        test_slf_img = test_slf_img.reshape(-1, self.args.K0 * self.args.K1)

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab = testing_data['ab_norm'][idx]

        test_input_RSS = torch.from_numpy(test_input_RSS).float()
        test_denoise_RSS = torch.from_numpy(test_denoise_RSS).float()
        test_slf_img = torch.from_numpy(test_slf_img).float()
        test_ab = torch.from_numpy(test_ab).float()
        test_noise_class = torch.from_numpy(test_noise_class).long()

        Test_Dataset = TensorDataset(test_input_RSS, test_denoise_RSS, test_slf_img, test_ab, test_noise_class)
        test_loader = DataLoader(
            Test_Dataset,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        print("Data Loaded for testing metrics!")
        return test_loader

    def test_model_ROC(self, noise_level='all'):
        model = MAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))
        model = model.to(self.device)
        model_path = self.path + 'model_' + str(self.args.n_epochs) + '.pth'
        model.load_state_dict(torch.load(model_path, map_location=self.device))
        model.eval()
        test_loader = self.load_data_metrics(noise_level)
        noise_class = 0  # noise_class = 0 stands for all noise classes
        if noise_level == 'low':
            noise_class = 1
        elif noise_level == 'mid':
            noise_class = 2
        elif noise_level == 'high':
            noise_class = 3

        # test model for testing set
        slf_rmse = 0
        slf_mae = 0
        labels = list()
        preds = list()
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(
                test_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            RSS, SLF, ab, noise_pred = model(input_RSS)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            labels.append(target_SLF.detach().cpu().numpy())
            preds.append(SLF.detach().cpu().numpy())

        slf_rmse /= len(test_loader.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader.dataset)
        print('====> Test for Testing Set:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))

        preds = np.concatenate(preds, axis=0)
        np.save('results/MAE/ROC_data/mae_preds_' + noise_level + '.npy', preds)

    def estimate(self, noise_level='low'):
        if noise_level != 'low' and noise_level != 'mid' and noise_level != 'high' :
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))
        model = MAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))
        model = model.to(self.device)
        model_path = self.path + 'model_' + str(self.args.n_epochs) + '.pth'
        model.load_state_dict(torch.load(model_path, map_location=self.device))
        model.eval()
        test_loader = self.load_data_metrics(noise_level)
        path_model = self.path + 'YOLO/Origin/images/mae/' + noise_level + '/'
        if not os.path.exists(path_model):
            os.makedirs(path_model)
        print('====> Test Model with estimation:')
        print('====> Noise Level: ' + noise_level)
        thr = 0.4
        # test model
        for batch_idx, (input_rss, _, target_slf, *_) in enumerate(test_loader):
            input_rss = input_rss.to(self.device)
            target_slf = target_slf.to(self.device)
            _, slf_recon, *_ = model(input_rss)
            slf_recon = slf_recon.detach().cpu().numpy()
            for idx in range(len(target_slf)):
                estimator_img = slf_recon[idx, :].reshape(self.args.K0, self.args.K1)
                # estimator_img[estimator_img > thr] = 1
                # estimator_img[estimator_img <= thr] = 0
                estimator_img = cv2.resize(estimator_img, (320, 320), interpolation=cv2.INTER_NEAREST)*255
                cv2.imwrite(path_model + 'im' + str(idx+batch_idx*self.args.batch_size+1) +'.jpg', estimator_img)
            image = cv2.imread(path_model + 'im' + str(1) +'.jpg', cv2.IMREAD_GRAYSCALE)


# runner for Multi-task Gaussian VAE
class MGVAE_runner():
    def __init__(self, args):
        self.args = args
        cuda_flag = not self.args.no_cuda and torch.cuda.is_available()
        self.device = torch.device("cuda" if cuda_flag else "cpu")
        torch.manual_seed(self.args.seed)
        random.seed(self.args.seed)
        self.path = self.args.run + 'MGVAE/'
        if not os.path.exists(self.path):
            os.makedirs(self.path)

    def get_optimizer(self, parameters):
        if self.args.optimizer == 'Adam':
            return optim.Adam(parameters, lr=self.args.lr, weight_decay=self.args.weight_decay, betas=(0.9, 0.999), eps=1e-7)
        elif self.args.optimizer == 'RMSProp':
            return optim.RMSprop(parameters, lr=self.args.lr, weight_decay=self.args.weight_decay)
        elif self.args.optimizer == 'SGD':
            return optim.SGD(parameters, lr=self.args.lr, momentum=0.9)
        else:
            raise NotImplementedError('Optimizer {} not understood.'.format(self.args.optimizer))

    def load_data(self, noise_level='all'):
        dataset, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min = self.load_training_data()
        train_size = int(len(dataset) * 0.9)
        val_size = int(len(dataset) * 0.1)
        dataset_train, dataset_val = random_split(dataset, [train_size, val_size],
                                                  generator=torch.Generator().manual_seed(self.args.seed))
        dataset_test = self.load_testing_data(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        dataset_test1 = self.load_testing_img1(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        dataset_test2 = self.load_testing_img2(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        train_loader = DataLoader(
            dataset_train,
            batch_size=self.args.batch_size, shuffle=True, num_workers=2)
        val_loader = DataLoader(
            dataset_val,
            batch_size=self.args.batch_size, shuffle=True, num_workers=2)
        test_loader = DataLoader(
            dataset_test,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        test_loader1 = DataLoader(
            dataset_test1,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        test_loader2 = DataLoader(
            dataset_test2,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        print("Data Loaded!")
        return train_loader, val_loader, test_loader, test_loader1, test_loader2

    def load_training_data(self):
        # Load training data
        training_data = sio.loadmat('./data/training_data.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        train_input_RSS = training_data['RSS']
        RSS_max = np.amax(train_input_RSS, axis=0, keepdims=True)
        RSS_min = np.amin(train_input_RSS, axis=0, keepdims=True)
        train_input_RSS = (train_input_RSS - RSS_min) / (RSS_max - RSS_min)
        train_input_RSS = train_input_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        train_denoise_RSS = training_data['RSS_denoise']
        RSS_denoise_max = np.amax(train_denoise_RSS, axis=0, keepdims=True)
        RSS_denoise_min = np.amin(train_denoise_RSS, axis=0, keepdims=True)
        train_denoise_RSS = (train_denoise_RSS - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        train_denoise_RSS = train_denoise_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        train_slf_img = training_data['slf_img']
        train_slf_img = train_slf_img.reshape(-1, self.args.K0 * self.args.K1)
        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        train_ab = training_data['ab_norm']
        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        train_noise_level = training_data['sig_epsilon_class']
        train_noise_level = np.squeeze(train_noise_level)

        train_input_RSS = torch.from_numpy(train_input_RSS).float()
        train_denoise_RSS = torch.from_numpy(train_denoise_RSS).float()
        train_slf_img = torch.from_numpy(train_slf_img).float()
        train_ab = torch.from_numpy(train_ab).float()
        train_noise_level = torch.from_numpy(train_noise_level).long()

        Train_Dataset = TensorDataset(train_input_RSS, train_denoise_RSS, train_slf_img, train_ab, train_noise_level)
        return Train_Dataset, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min

    def load_testing_data(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data
        testing_data = sio.loadmat('./data/testing_data.mat')
        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_class = testing_data['sig_epsilon_class']
        test_noise_class = np.squeeze(test_noise_class)
        idx = (test_noise_class == 0) | (test_noise_class == 1) | (test_noise_class == 2)
        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            idx = (test_noise_class == 0)
        elif noise_level == 'mid':
            idx = (test_noise_class == 1)
        elif noise_level == 'high':
            idx = (test_noise_class == 2)
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))
        test_noise_class = test_noise_class[idx]
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS = testing_data['RSS'][idx]
        test_input_RSS = (test_input_RSS - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS = test_input_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS = testing_data['RSS_denoise'][idx]
        test_denoise_RSS = (test_denoise_RSS - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS = test_denoise_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img = testing_data['slf_img'][idx]
        test_slf_img = test_slf_img.reshape(-1, self.args.K0 * self.args.K1)

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab = testing_data['ab_norm'][idx]

        test_input_RSS = torch.from_numpy(test_input_RSS).float()
        test_denoise_RSS = torch.from_numpy(test_denoise_RSS).float()
        test_slf_img = torch.from_numpy(test_slf_img).float()
        test_ab = torch.from_numpy(test_ab).float()
        test_noise_class = torch.from_numpy(test_noise_class).long()

        Test_Dataset = TensorDataset(test_input_RSS, test_denoise_RSS, test_slf_img, test_ab, test_noise_class)
        return Test_Dataset

    def load_testing_img1(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data for image 1
        testing_data1 = sio.loadmat('./data/testing_data_image1.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS1 = testing_data1['RSS']
        test_input_RSS1 = (test_input_RSS1 - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS1 = test_input_RSS1.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS1 = testing_data1['RSS_denoise']
        test_denoise_RSS1 = (test_denoise_RSS1 - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS1 = test_denoise_RSS1.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img1 = testing_data1['slf_img']
        test_slf_img1 = test_slf_img1.reshape(1, self.args.K0 * self.args.K1)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        test_slf_img1 = np.tile(test_slf_img1, (test_input_RSS1.shape[0], 1))

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab1 = testing_data1['ab_norm']

        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_level1 = testing_data1['sig_epsilon_class']
        test_noise_level1 = np.squeeze(test_noise_level1)

        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            test_input_RSS1 = test_input_RSS1[:500]
            test_denoise_RSS1 = test_denoise_RSS1[:500]
            test_slf_img1 = test_slf_img1[:500]
            test_ab1 = test_ab1[:500]
            test_noise_level1 = test_noise_level1[:500]
        elif noise_level == 'mid':
            test_input_RSS1 = test_input_RSS1[500:1000]
            test_denoise_RSS1 = test_denoise_RSS1[500:1000]
            test_slf_img1 = test_slf_img1[500:1000]
            test_ab1 = test_ab1[500:1000]
            test_noise_level1 = test_noise_level1[500:1000]
        elif noise_level == 'high':
            test_input_RSS1 = test_input_RSS1[1000:]
            test_denoise_RSS1 = test_denoise_RSS1[1000:]
            test_slf_img1 = test_slf_img1[1000:]
            test_ab1 = test_ab1[1000:]
            test_noise_level1 = test_noise_level1[1000:]
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))

        test_input_RSS1 = torch.from_numpy(test_input_RSS1).float()
        test_denoise_RSS1 = torch.from_numpy(test_denoise_RSS1).float()
        test_slf_img1 = torch.from_numpy(test_slf_img1).float()
        test_ab1 = torch.from_numpy(test_ab1).float()
        test_noise_level1 = torch.from_numpy(test_noise_level1).long()

        Test_Dataset_1 = TensorDataset(test_input_RSS1, test_denoise_RSS1, test_slf_img1, test_ab1, test_noise_level1)
        return Test_Dataset_1

    def load_testing_img2(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data for image 2
        testing_data2 = sio.loadmat('./data/testing_data_image2.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS2 = testing_data2['RSS']
        test_input_RSS2 = (test_input_RSS2 - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS2 = test_input_RSS2.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS2 = testing_data2['RSS_denoise_norm']
        test_denoise_RSS2 = (test_denoise_RSS2 - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS2 = test_denoise_RSS2.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img2 = testing_data2['slf_img']
        test_slf_img2 = test_slf_img2.reshape(1, self.args.K0 * self.args.K1)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        test_slf_img2 = np.tile(test_slf_img2, (test_input_RSS2.shape[0], 1))

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab2 = testing_data2['ab_norm']

        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_level2 = testing_data2['sig_epsilon_class']
        test_noise_level2 = np.squeeze(test_noise_level2)

        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            test_input_RSS2 = test_input_RSS2[:500]
            test_denoise_RSS2 = test_denoise_RSS2[:500]
            test_slf_img2 = test_slf_img2[:500]
            test_ab2 = test_ab2[:500]
            test_noise_level2 = test_noise_level2[:500]
        elif noise_level == 'mid':
            test_input_RSS2 = test_input_RSS2[500:1000]
            test_denoise_RSS2 = test_denoise_RSS2[500:1000]
            test_slf_img2 = test_slf_img2[500:1000]
            test_ab2 = test_ab2[500:1000]
            test_noise_level2 = test_noise_level2[500:1000]
        elif noise_level == 'high':
            test_input_RSS2 = test_input_RSS2[1000:]
            test_denoise_RSS2 = test_denoise_RSS2[1000:]
            test_slf_img2 = test_slf_img2[1000:]
            test_ab2 = test_ab2[1000:]
            test_noise_level2 = test_noise_level2[1000:]
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))

        test_input_RSS2 = torch.from_numpy(test_input_RSS2).float()
        test_denoise_RSS2 = torch.from_numpy(test_denoise_RSS2).float()
        test_slf_img2 = torch.from_numpy(test_slf_img2).float()
        test_ab2 = torch.from_numpy(test_ab2).float()
        test_noise_level2 = torch.from_numpy(test_noise_level2).long()

        Test_Dataset_2 = TensorDataset(test_input_RSS2, test_denoise_RSS2, test_slf_img2, test_ab2, test_noise_level2)
        return Test_Dataset_2

    def train(self, model, train_loader, optimizer, epoch):
        model.train()
        train_loss = 0                  # total loss
        train_kld = 0                   # kl-divergence
        train_loss_task0 = 0            # loss for reconstruct RSS
        train_loss_task1 = 0            # loss for SLF image estimation
        train_loss_task2 = 0            # loss for parameters estimation
        train_loss_task3 = 0            # loss for noise level prediction
        train_acc = 0                   # accuracy for noise level prediction
        batch_num = int(len(train_loader.dataset) / self.args.batch_size) + 1
        batch_10num = batch_num * 10
        for batch_idx,  (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(train_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            optimizer.zero_grad()
            KLD, RSS, SLF, ab, noise_level = model(input_RSS)
            loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum')/(self.args.N * self.args.P * self.args.P)
            loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum')/(self.args.K0 * self.args.K1)
            loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum')/(self.args.N + 1)
            loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
            loss_weight0 = (((epoch-1) % 20) * batch_num + batch_idx) / batch_10num
            loss_weight0 = np.minimum(loss_weight0, self.args.lambda0)
            loss = loss_weight0 * KLD + self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1\
                   + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3
            train_loss += loss.item()
            train_kld += KLD.item()
            train_loss_task0 += loss_task0.item()
            train_loss_task1 += loss_task1.item()
            train_loss_task2 += loss_task2.item()
            train_loss_task3 += loss_task3.item()
            _, noise_pred = torch.max(noise_level, 1)
            acc_task3 = (noise_pred == target_noise_level).sum()
            train_acc += acc_task3.item()
            loss /= input_RSS.size(0)
            loss.backward()
            optimizer.step()

            if batch_idx % self.args.log_interval == 0:
                print(
                    'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tKLD: {:.6f}\tLoss0: {:.6f}\tLoss1: {:.6f}\tLoss2: {:.6f}\tLoss3: {:.6f}\tAcc: {:.6f}'.format(
                        epoch, batch_idx * len(input_RSS), len(train_loader.dataset),
                               100. * batch_idx / len(train_loader),
                               loss.item(),
                               KLD.item() / len(input_RSS),
                               loss_task0.item() / len(input_RSS), loss_task1.item() / len(input_RSS),
                               loss_task2.item() / len(input_RSS), loss_task3.item() / len(input_RSS), acc_task3.item() / len(input_RSS)))

        train_loss /= len(train_loader.dataset)
        train_kld /= len(train_loader.dataset)
        train_loss_task0 /= len(train_loader.dataset)
        train_loss_task1 /= len(train_loader.dataset)
        train_loss_task2 /= len(train_loader.dataset)
        train_loss_task3 /= len(train_loader.dataset)
        train_acc /= len(train_loader.dataset)
        print('====> Epoch: {} Average loss: {:.6f}'.format(epoch, train_loss))
        print('====> KL divergence: {:.6f}'.format(train_kld))
        print('====> RSS loss: {:.6f}'.format(train_loss_task0))
        print('====> SLF loss: {:.6f}'.format(train_loss_task1))
        print('====> ab loss: {:.6f}'.format(train_loss_task2))
        print('====> noise level loss: {:.6f}'.format(train_loss_task3))
        print('====> noise level Accuracy: {:.4f}'.format(train_acc))
        return train_loss, train_kld, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3, train_acc

    def validate(self, model, val_loader, epoch):
        model.eval()
        val_loss = 0                  # total loss
        val_rmse_slf = 0  # slf reconstruction loss
        val_mae_slf = 0  # slf reconstruction loss
        batch_num = int(len(val_loader.dataset) / self.args.batch_size) + 1
        batch_10num = batch_num * 10
        for batch_idx,  (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(val_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_level = model(input_RSS)
            loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum')/(self.args.N * self.args.P * self.args.P)
            loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum')/(self.args.K0 * self.args.K1)
            loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum')/(self.args.N + 1)
            loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
            loss_weight0 = (((epoch - 1) % 20) * batch_num + batch_idx) / batch_10num
            loss_weight0 = np.minimum(loss_weight0, self.args.lambda0)
            loss = 0 * KLD + self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1 \
                   + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3
            val_loss += loss.item()
            slf_mse = F.mse_loss(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
            slf_mae = F.l1_loss(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
            val_rmse_slf += slf_mse.item()
            val_mae_slf += slf_mae.item()
        val_loss /= len(val_loader.dataset)
        val_rmse_slf /= len(val_loader.dataset)
        val_rmse_slf = np.sqrt(val_rmse_slf)
        val_mae_slf /= len(val_loader.dataset)
        print('====> Validation set loss: {:.6f}'.format(val_loss))
        print('====> Validation set SLF rmse: {:.6f}'.format(val_rmse_slf))
        print('====> Validation set SLF mae: {:.6f}'.format(val_mae_slf))
        return val_loss, val_rmse_slf, val_mae_slf

    def test(self, model, test_loader, epoch):
        model.eval()
        test_loss = 0  # total loss
        test_kld = 0   # KL-divergence
        test_loss_task0 = 0  # loss for reconstruct RSS
        test_loss_task1 = 0  # loss for SLF image estimation
        test_loss_task2 = 0  # loss for parameters estimation
        test_loss_task3 = 0  # loss for noise level prediction
        batch_num = int(len(test_loader.dataset) / self.args.batch_size) + 1
        batch_10num = batch_num * 10
        with torch.no_grad():
            for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader):
                input_RSS = input_RSS.to(self.device)
                target_RSS = target_RSS.to(self.device)
                target_SLF = target_SLF.to(self.device)
                target_ab = target_ab.to(self.device)
                target_noise_level = target_noise_level.to(self.device)
                KLD, RSS, SLF, ab, noise_level = model(input_RSS)
                loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum') / (
                            self.args.N * self.args.P * self.args.P)
                loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
                loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum') / (self.args.N + 1)
                loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
                loss_weight0 = (((epoch - 1) % 20) * batch_num + batch_idx) / batch_10num
                loss_weight0 = np.minimum(loss_weight0, self.args.lambda0)
                loss = loss_weight0 * KLD + self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1 \
                       + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3

                test_loss += loss.item()
                test_kld += KLD.item()
                test_loss_task0 += loss_task0.item()
                test_loss_task1 += loss_task1.item()
                test_loss_task2 += loss_task2.item()
                test_loss_task3 += loss_task3.item()

        test_loss /= len(test_loader.dataset)
        test_kld /= len(test_loader.dataset)
        test_loss_task0 /= len(test_loader.dataset)
        test_loss_task1 /= len(test_loader.dataset)
        test_loss_task2 /= len(test_loader.dataset)
        test_loss_task3 /= len(test_loader.dataset)
        test_acc = self.test_accuracy(model, test_loader)
        print('====> Test set loss: {:.6f}'.format(test_loss))
        print('====> Test set KL divergence: {:.6f}'.format(test_kld))
        print('====> Test set RSS loss: {:.6f}'.format(test_loss_task0))
        print('====> Test set SLF loss: {:.6f}'.format(test_loss_task1))
        print('====> Test set ab loss: {:.6f}'.format(test_loss_task2))
        print('====> Test set noise level loss: {:.6f}'.format(test_loss_task3))
        print('====> Test set noise level Accuracy: {:.4f}'.format(test_acc))
        return test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc

    def train_save(self):
        model = MGVAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))

        def weights_init(m):
            if isinstance(m, torch.nn.Linear) or isinstance(m, torch.nn.Conv2d) or isinstance(m, torch.nn.ConvTranspose2d):
                torch.nn.init.xavier_uniform_(m.weight.data)
                if m.bias is not None:
                    torch.nn.init.zeros_(m.bias.data)

        model.apply(weights_init)
        model = model.to(self.device)
        summary(model, input_size=(self.args.batch_size, self.args.N, self.args.P, self.args.P))
        train_loader, val_loader, test_loader, test_loader1, test_loader2 = self.load_data()
        optimizer = self.get_optimizer(model.parameters())
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.25, patience=10, min_lr=1e-5)
        train_set = []
        val_set = []
        test_set = []       # for test set
        test_set1 = []      # for test image 1
        test_set2 = []      # for test image 2
        train_loss_path = self.path + 'train_loss_' + str(self.args.n_epochs) + '.npy'
        val_loss_path = self.path + 'val_loss_' + str(self.args.n_epochs) + '.npy'
        test_loss_path = self.path + 'test_loss_' + str(self.args.n_epochs) + '.npz'
        model_path = self.path + 'model_' + str(self.args.n_epochs) + '.pth'
        current_val_slf_loss = np.inf
        for epoch in range(1, self.args.n_epochs+1):
            print("Epoch %d learning rate：%f" % (epoch, optimizer.param_groups[0]['lr']))
            train_loss, train_kld, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3, train_acc = self.train(model,
                                                                                                            train_loader,
                                                                                                            optimizer,
                                                                                                            epoch)
            train_set.append([train_loss, train_kld, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3, train_acc])
            print('====> Validation Loss for Epoch {:d}'.format(epoch))
            val_loss, val_slf_rmse, val_slf_mae = self.validate(model, val_loader, epoch)
            val_set.append([val_loss, val_slf_rmse, val_slf_mae])
            scheduler.step(val_loss)
            print('====> Test for Epoch {:d}'.format(epoch))
            test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc = self.test(
                model,
                test_loader, epoch)
            test_set.append(
                [test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc])
            print('====> Test Image1 for Epoch {:d}'.format(epoch))
            test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc = self.test(model,
                                                                                                                test_loader1, epoch)
            test_set1.append([test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc])
            print('====> Test Image2 for Epoch {:d}'.format(epoch))
            test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc = self.test(model,
                                                                                                                test_loader2, epoch)
            test_set2.append([test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc])
            if epoch % self.args.save_freq == 0:
                ckpt_path = self.path + '/model' + '_' + str(epoch) + '.pth'
                torch.save(model.state_dict(), ckpt_path)
                print('checkpoint{}.pth saved!'.format(epoch))
            if val_slf_mae < current_val_slf_loss:
                torch.save(model.state_dict(), model_path)
                current_val_slf_loss = val_slf_mae

        train_loss = np.asarray(train_set).reshape(-1, 7)
        val_loss = np.asarray(val_set).reshape(-1, 3)
        test_loss = np.asarray(test_set).reshape(-1, 7)
        test_loss1 = np.asarray(test_set1).reshape(-1, 7)
        test_loss2 = np.asarray(test_set2).reshape(-1, 7)
        np.save(train_loss_path, train_loss)
        np.save(val_loss_path, val_loss)
        np.savez(test_loss_path, test_loss=test_loss, test_loss1=test_loss1, test_loss2=test_loss2)

    def test_accuracy(self, model, test_loader):
        model.eval()
        acc = 0
        for i, (input_RSS, *_, target_noise_level) in enumerate(test_loader):
            input_RSS = input_RSS.to(self.device)
            *_, noise_pred = model(input_RSS)  # noise_pred: [batch_size, 3]
            noise_pred = noise_pred.detach().cpu().numpy()
            noise_pred = np.argmax(noise_pred, axis=1)
            labels = target_noise_level.detach().cpu().numpy()
            acc += np.sum((noise_pred == labels).astype(int))
        acc /= len(test_loader.dataset)
        return acc

    def test_model(self, noise_level='all'):
        model = MGVAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))
        model = model.to(self.device)
        model_path = self.path + 'model_' + str(self.args.n_epochs) + '.pth'
        model.load_state_dict(torch.load(model_path, map_location=self.device))
        model.eval()
        *_, test_loader, test_loader1, test_loader2 = self.load_data(noise_level)
        noise_class = 0  # noise_class = 0 stands for all noise classes
        if noise_level == 'low':
            noise_class = 1
        elif noise_level == 'mid':
            noise_class = 2
        elif noise_level == 'high':
            noise_class = 3

        # test model for testing set
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(
                test_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_pred = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader.dataset)
        ab_rmse /= len(test_loader.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader)
        print('====> Test for Testing Set:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))

        # test model for image 1
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader1):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_pred = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF1_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader1.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader1.dataset)
        ab_rmse /= len(test_loader1.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader1)
        print('====> Test for image 1:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))

        # test model for image 2
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader2):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_pred = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF2_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader2.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader2.dataset)
        ab_rmse /= len(test_loader2.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader2)
        print('====> Test for image 2:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))


# runner for Multi-task Laplacian VAE
class MLVAE_runner():
    def __init__(self, args):
        self.args = args
        cuda_flag = not self.args.no_cuda and torch.cuda.is_available()
        self.device = torch.device("cuda" if cuda_flag else "cpu")
        torch.manual_seed(self.args.seed)
        random.seed(self.args.seed)
        self.path = self.args.run + 'MLVAE/'
        if not os.path.exists(self.path):
            os.makedirs(self.path)

    def get_optimizer(self, parameters):
        if self.args.optimizer == 'Adam':
            return optim.Adam(parameters, lr=self.args.lr, weight_decay=self.args.weight_decay, betas=(0.9, 0.999), eps=1e-7)
        elif self.args.optimizer == 'RMSProp':
            return optim.RMSprop(parameters, lr=self.args.lr, weight_decay=self.args.weight_decay)
        elif self.args.optimizer == 'SGD':
            return optim.SGD(parameters, lr=self.args.lr, momentum=0.9)
        else:
            raise NotImplementedError('Optimizer {} not understood.'.format(self.args.optimizer))

    def load_data(self, noise_level='all'):
        dataset, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min = self.load_training_data()
        train_size = int(len(dataset) * 0.9)
        val_size = int(len(dataset) * 0.1)
        dataset_train, dataset_val = random_split(dataset, [train_size, val_size],
                                                  generator=torch.Generator().manual_seed(self.args.seed))
        dataset_test = self.load_testing_data(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        dataset_test1 = self.load_testing_img1(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        dataset_test2 = self.load_testing_img2(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        train_loader = DataLoader(
            dataset_train,
            batch_size=self.args.batch_size, shuffle=True, num_workers=2)
        val_loader = DataLoader(
            dataset_val,
            batch_size=self.args.batch_size, shuffle=True, num_workers=2)
        test_loader = DataLoader(
            dataset_test,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        test_loader1 = DataLoader(
            dataset_test1,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        test_loader2 = DataLoader(
            dataset_test2,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        print("Data Loaded!")
        return train_loader, val_loader, test_loader, test_loader1, test_loader2

    def load_training_data(self):
        # Load training data
        training_data = sio.loadmat('./data/training_data.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        train_input_RSS = training_data['RSS']
        RSS_max = np.amax(train_input_RSS, axis=0, keepdims=True)
        RSS_min = np.amin(train_input_RSS, axis=0, keepdims=True)
        train_input_RSS = (train_input_RSS - RSS_min) / (RSS_max - RSS_min)
        train_input_RSS = train_input_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        train_denoise_RSS = training_data['RSS_denoise']
        RSS_denoise_max = np.amax(train_denoise_RSS, axis=0, keepdims=True)
        RSS_denoise_min = np.amin(train_denoise_RSS, axis=0, keepdims=True)
        train_denoise_RSS = (train_denoise_RSS - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        train_denoise_RSS = train_denoise_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        train_slf_img = training_data['slf_img']
        train_slf_img = train_slf_img.reshape(-1, self.args.K0 * self.args.K1)
        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        train_ab = training_data['ab_norm']
        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        train_noise_level = training_data['sig_epsilon_class']
        train_noise_level = np.squeeze(train_noise_level)

        train_input_RSS = torch.from_numpy(train_input_RSS).float()
        train_denoise_RSS = torch.from_numpy(train_denoise_RSS).float()
        train_slf_img = torch.from_numpy(train_slf_img).float()
        train_ab = torch.from_numpy(train_ab).float()
        train_noise_level = torch.from_numpy(train_noise_level).long()

        Train_Dataset = TensorDataset(train_input_RSS, train_denoise_RSS, train_slf_img, train_ab, train_noise_level)
        return Train_Dataset, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min

    def load_testing_data(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data
        testing_data = sio.loadmat('./data/testing_data.mat')
        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_class = testing_data['sig_epsilon_class']
        test_noise_class = np.squeeze(test_noise_class)
        idx = (test_noise_class == 0) | (test_noise_class == 1) | (test_noise_class == 2)
        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            idx = (test_noise_class == 0)
        elif noise_level == 'mid':
            idx = (test_noise_class == 1)
        elif noise_level == 'high':
            idx = (test_noise_class == 2)
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))
        test_noise_class = test_noise_class[idx]
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS = testing_data['RSS'][idx]
        test_input_RSS = (test_input_RSS - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS = test_input_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS = testing_data['RSS_denoise'][idx]
        test_denoise_RSS = (test_denoise_RSS - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS = test_denoise_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img = testing_data['slf_img'][idx]
        test_slf_img = test_slf_img.reshape(-1, self.args.K0 * self.args.K1)

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab = testing_data['ab_norm'][idx]

        test_input_RSS = torch.from_numpy(test_input_RSS).float()
        test_denoise_RSS = torch.from_numpy(test_denoise_RSS).float()
        test_slf_img = torch.from_numpy(test_slf_img).float()
        test_ab = torch.from_numpy(test_ab).float()
        test_noise_class = torch.from_numpy(test_noise_class).long()

        Test_Dataset = TensorDataset(test_input_RSS, test_denoise_RSS, test_slf_img, test_ab, test_noise_class)
        return Test_Dataset

    def load_testing_img1(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data for image 1
        testing_data1 = sio.loadmat('./data/testing_data_image1.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS1 = testing_data1['RSS']
        test_input_RSS1 = (test_input_RSS1 - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS1 = test_input_RSS1.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS1 = testing_data1['RSS_denoise']
        test_denoise_RSS1 = (test_denoise_RSS1 - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS1 = test_denoise_RSS1.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img1 = testing_data1['slf_img']
        test_slf_img1 = test_slf_img1.reshape(1, self.args.K0 * self.args.K1)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        test_slf_img1 = np.tile(test_slf_img1, (test_input_RSS1.shape[0], 1))

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab1 = testing_data1['ab_norm']

        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_level1 = testing_data1['sig_epsilon_class']
        test_noise_level1 = np.squeeze(test_noise_level1)

        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            test_input_RSS1 = test_input_RSS1[:500]
            test_denoise_RSS1 = test_denoise_RSS1[:500]
            test_slf_img1 = test_slf_img1[:500]
            test_ab1 = test_ab1[:500]
            test_noise_level1 = test_noise_level1[:500]
        elif noise_level == 'mid':
            test_input_RSS1 = test_input_RSS1[500:1000]
            test_denoise_RSS1 = test_denoise_RSS1[500:1000]
            test_slf_img1 = test_slf_img1[500:1000]
            test_ab1 = test_ab1[500:1000]
            test_noise_level1 = test_noise_level1[500:1000]
        elif noise_level == 'high':
            test_input_RSS1 = test_input_RSS1[1000:]
            test_denoise_RSS1 = test_denoise_RSS1[1000:]
            test_slf_img1 = test_slf_img1[1000:]
            test_ab1 = test_ab1[1000:]
            test_noise_level1 = test_noise_level1[1000:]
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))

        test_input_RSS1 = torch.from_numpy(test_input_RSS1).float()
        test_denoise_RSS1 = torch.from_numpy(test_denoise_RSS1).float()
        test_slf_img1 = torch.from_numpy(test_slf_img1).float()
        test_ab1 = torch.from_numpy(test_ab1).float()
        test_noise_level1 = torch.from_numpy(test_noise_level1).long()

        Test_Dataset_1 = TensorDataset(test_input_RSS1, test_denoise_RSS1, test_slf_img1, test_ab1, test_noise_level1)
        return Test_Dataset_1

    def load_testing_img2(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data for image 2
        testing_data2 = sio.loadmat('./data/testing_data_image2.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS2 = testing_data2['RSS']
        test_input_RSS2 = (test_input_RSS2 - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS2 = test_input_RSS2.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS2 = testing_data2['RSS_denoise_norm']
        test_denoise_RSS2 = (test_denoise_RSS2 - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS2 = test_denoise_RSS2.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img2 = testing_data2['slf_img']
        test_slf_img2 = test_slf_img2.reshape(1, self.args.K0 * self.args.K1)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        test_slf_img2 = np.tile(test_slf_img2, (test_input_RSS2.shape[0], 1))

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab2 = testing_data2['ab_norm']

        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_level2 = testing_data2['sig_epsilon_class']
        test_noise_level2 = np.squeeze(test_noise_level2)

        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            test_input_RSS2 = test_input_RSS2[:500]
            test_denoise_RSS2 = test_denoise_RSS2[:500]
            test_slf_img2 = test_slf_img2[:500]
            test_ab2 = test_ab2[:500]
            test_noise_level2 = test_noise_level2[:500]
        elif noise_level == 'mid':
            test_input_RSS2 = test_input_RSS2[500:1000]
            test_denoise_RSS2 = test_denoise_RSS2[500:1000]
            test_slf_img2 = test_slf_img2[500:1000]
            test_ab2 = test_ab2[500:1000]
            test_noise_level2 = test_noise_level2[500:1000]
        elif noise_level == 'high':
            test_input_RSS2 = test_input_RSS2[1000:]
            test_denoise_RSS2 = test_denoise_RSS2[1000:]
            test_slf_img2 = test_slf_img2[1000:]
            test_ab2 = test_ab2[1000:]
            test_noise_level2 = test_noise_level2[1000:]
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))

        test_input_RSS2 = torch.from_numpy(test_input_RSS2).float()
        test_denoise_RSS2 = torch.from_numpy(test_denoise_RSS2).float()
        test_slf_img2 = torch.from_numpy(test_slf_img2).float()
        test_ab2 = torch.from_numpy(test_ab2).float()
        test_noise_level2 = torch.from_numpy(test_noise_level2).long()

        Test_Dataset_2 = TensorDataset(test_input_RSS2, test_denoise_RSS2, test_slf_img2, test_ab2, test_noise_level2)
        return Test_Dataset_2

    def train(self, model, train_loader, optimizer, epoch):
        model.train()
        train_loss = 0  # total loss
        train_kld = 0  # kl-divergence
        train_loss_task0 = 0  # loss for reconstruct RSS
        train_loss_task1 = 0  # loss for SLF image estimation
        train_loss_task2 = 0  # loss for parameters estimation
        train_loss_task3 = 0  # loss for noise level prediction
        train_acc = 0  # accuracy for noise level prediction
        batch_num = int(len(train_loader.dataset) / self.args.batch_size) + 1
        batch_10num = batch_num * 10
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(train_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            optimizer.zero_grad()
            KLD, RSS, SLF, ab, noise_level = model(input_RSS)
            loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum') / (
                        self.args.N * self.args.P * self.args.P)
            loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
            loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum') / (self.args.N + 1)
            loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
            loss_weight0 = (((epoch - 1) % 20) * batch_num + batch_idx) / batch_10num
            loss_weight0 = np.minimum(loss_weight0, self.args.lambda0)
            loss = loss_weight0 * KLD + self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1 \
                   + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3
            train_loss += loss.item()
            train_kld += KLD.item()
            train_loss_task0 += loss_task0.item()
            train_loss_task1 += loss_task1.item()
            train_loss_task2 += loss_task2.item()
            train_loss_task3 += loss_task3.item()
            _, noise_pred = torch.max(noise_level, 1)
            acc_task3 = (noise_pred == target_noise_level).sum()
            train_acc += acc_task3.item()
            loss /= input_RSS.size(0)
            loss.backward()
            optimizer.step()

            if batch_idx % self.args.log_interval == 0:
                print(
                    'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tKLD: {:.6f}\tLoss0: {:.6f}\tLoss1: {:.6f}\tLoss2: {:.6f}\tLoss3: {:.6f}\tAcc: {:.6f}'.format(
                        epoch, batch_idx * len(input_RSS), len(train_loader.dataset),
                               100. * batch_idx / len(train_loader),
                        loss.item(),
                               KLD.item() / len(input_RSS),
                               loss_task0.item() / len(input_RSS), loss_task1.item() / len(input_RSS),
                               loss_task2.item() / len(input_RSS), loss_task3.item() / len(input_RSS),
                               acc_task3.item() / len(input_RSS)))

        train_loss /= len(train_loader.dataset)
        train_kld /= len(train_loader.dataset)
        train_loss_task0 /= len(train_loader.dataset)
        train_loss_task1 /= len(train_loader.dataset)
        train_loss_task2 /= len(train_loader.dataset)
        train_loss_task3 /= len(train_loader.dataset)
        train_acc /= len(train_loader.dataset)
        print('====> Epoch: {} Average loss: {:.6f}'.format(epoch, train_loss))
        print('====> KL divergence: {:.6f}'.format(train_kld))
        print('====> RSS loss: {:.6f}'.format(train_loss_task0))
        print('====> SLF loss: {:.6f}'.format(train_loss_task1))
        print('====> ab loss: {:.6f}'.format(train_loss_task2))
        print('====> noise level loss: {:.6f}'.format(train_loss_task3))
        print('====> noise level Accuracy: {:.4f}'.format(train_acc))
        return train_loss, train_kld, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3, train_acc

    def validate(self, model, val_loader, epoch):
        model.eval()
        val_loss = 0  # total loss
        val_rmse_slf = 0  # slf reconstruction loss
        val_mae_slf = 0  # slf reconstruction loss
        batch_num = int(len(val_loader.dataset) / self.args.batch_size) + 1
        batch_10num = batch_num * 10
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(val_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_level = model(input_RSS)
            loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum') / (
                        self.args.N * self.args.P * self.args.P)
            loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
            loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum') / (self.args.N + 1)
            loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
            loss_weight0 = (((epoch - 1) % 20) * batch_num + batch_idx) / batch_10num
            loss_weight0 = np.minimum(loss_weight0, self.args.lambda0)
            loss = 0 * KLD + self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1 \
                   + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3
            val_loss += loss.item()
            slf_mse = F.mse_loss(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
            slf_mae = F.l1_loss(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
            val_rmse_slf += slf_mse.item()
            val_mae_slf += slf_mae.item()
        val_loss /= len(val_loader.dataset)
        val_rmse_slf /= len(val_loader.dataset)
        val_rmse_slf = np.sqrt(val_rmse_slf)
        val_mae_slf /= len(val_loader.dataset)
        print('====> Validation set loss: {:.6f}'.format(val_loss))
        print('====> Validation set SLF rmse: {:.6f}'.format(val_rmse_slf))
        print('====> Validation set SLF mae: {:.6f}'.format(val_mae_slf))
        return val_loss, val_rmse_slf, val_mae_slf

    def test(self, model, test_loader, epoch):
        model.eval()
        test_loss = 0  # total loss
        test_kld = 0  # KL-divergence
        test_loss_task0 = 0  # loss for reconstruct RSS
        test_loss_task1 = 0  # loss for SLF image estimation
        test_loss_task2 = 0  # loss for parameters estimation
        test_loss_task3 = 0  # loss for noise level prediction
        batch_num = int(len(test_loader.dataset) / self.args.batch_size) + 1
        batch_10num = batch_num * 10
        with torch.no_grad():
            for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader):
                input_RSS = input_RSS.to(self.device)
                target_RSS = target_RSS.to(self.device)
                target_SLF = target_SLF.to(self.device)
                target_ab = target_ab.to(self.device)
                target_noise_level = target_noise_level.to(self.device)
                KLD, RSS, SLF, ab, noise_level = model(input_RSS)
                loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum') / (
                        self.args.N * self.args.P * self.args.P)
                loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
                loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum') / (self.args.N + 1)
                loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
                loss_weight0 = (((epoch - 1) % 20) * batch_num + batch_idx) / batch_10num
                loss_weight0 = np.minimum(loss_weight0, self.args.lambda0)
                loss = loss_weight0 * KLD + self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1 \
                       + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3

                test_loss += loss.item()
                test_kld += KLD.item()
                test_loss_task0 += loss_task0.item()
                test_loss_task1 += loss_task1.item()
                test_loss_task2 += loss_task2.item()
                test_loss_task3 += loss_task3.item()

        test_loss /= len(test_loader.dataset)
        test_kld /= len(test_loader.dataset)
        test_loss_task0 /= len(test_loader.dataset)
        test_loss_task1 /= len(test_loader.dataset)
        test_loss_task2 /= len(test_loader.dataset)
        test_loss_task3 /= len(test_loader.dataset)
        test_acc = self.test_accuracy(model, test_loader)
        print('====> Test set loss: {:.6f}'.format(test_loss))
        print('====> Test set KL divergence: {:.6f}'.format(test_kld))
        print('====> Test set RSS loss: {:.6f}'.format(test_loss_task0))
        print('====> Test set SLF loss: {:.6f}'.format(test_loss_task1))
        print('====> Test set ab loss: {:.6f}'.format(test_loss_task2))
        print('====> Test set noise level loss: {:.6f}'.format(test_loss_task3))
        print('====> Test set noise level Accuracy: {:.4f}'.format(test_acc))
        return test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc

    def train_save(self):
        model = MLVAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))

        def weights_init(m):
            if isinstance(m, torch.nn.Linear) or isinstance(m, torch.nn.Conv2d) or isinstance(m,
                                                                                              torch.nn.ConvTranspose2d):
                torch.nn.init.xavier_uniform_(m.weight.data)
                if m.bias is not None:
                    torch.nn.init.zeros_(m.bias.data)

        model.apply(weights_init)
        model = model.to(self.device)
        summary(model, input_size=(self.args.batch_size, self.args.N, self.args.P, self.args.P))
        train_loader, val_loader, test_loader, test_loader1, test_loader2 = self.load_data()
        optimizer = self.get_optimizer(model.parameters())
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.25, patience=10, min_lr=1e-5)
        train_set = []
        val_set = []
        test_set = []  # for test set
        test_set1 = []  # for test image 1
        test_set2 = []  # for test image 2
        train_loss_path = self.path + 'train_loss_' + str(self.args.n_epochs) + '.npy'
        val_loss_path = self.path + 'val_loss_' + str(self.args.n_epochs) + '.npy'
        test_loss_path = self.path + 'test_loss_' + str(self.args.n_epochs) + '.npz'
        model_path = self.path + 'model_' + str(self.args.n_epochs) + '.pth'
        current_val_slf_loss = np.inf
        for epoch in range(1, self.args.n_epochs + 1):
            print("Epoch %d learning rate：%f" % (epoch, optimizer.param_groups[0]['lr']))
            train_loss, train_kld, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3, train_acc = self.train(
                model,
                train_loader,
                optimizer,
                epoch)
            train_set.append(
                [train_loss, train_kld, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3,
                 train_acc])
            print('====> Validation Loss for Epoch {:d}'.format(epoch))
            val_loss, val_slf_rmse, val_slf_mae = self.validate(model, val_loader, epoch)
            val_set.append([val_loss, val_slf_rmse, val_slf_mae])
            scheduler.step(val_loss)
            print('====> Test for Epoch {:d}'.format(epoch))
            test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc = self.test(
                model,
                test_loader, epoch)
            test_set.append(
                [test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc])
            print('====> Test Image1 for Epoch {:d}'.format(epoch))
            test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc = self.test(
                model,
                test_loader1, epoch)
            test_set1.append(
                [test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc])
            print('====> Test Image2 for Epoch {:d}'.format(epoch))
            test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc = self.test(
                model,
                test_loader2, epoch)
            test_set2.append(
                [test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc])
            if epoch % self.args.save_freq == 0:
                ckpt_path = self.path + '/model' + '_' + str(epoch) + '.pth'
                torch.save(model.state_dict(), ckpt_path)
                print('checkpoint{}.pth saved!'.format(epoch))
            if val_slf_mae < current_val_slf_loss:
                torch.save(model.state_dict(), model_path)
                current_val_slf_loss = val_slf_mae

        train_loss = np.asarray(train_set).reshape(-1, 7)
        val_loss = np.asarray(val_set).reshape(-1, 3)
        test_loss = np.asarray(test_set).reshape(-1, 7)
        test_loss1 = np.asarray(test_set1).reshape(-1, 7)
        test_loss2 = np.asarray(test_set2).reshape(-1, 7)
        np.save(train_loss_path, train_loss)
        np.save(val_loss_path, val_loss)
        np.savez(test_loss_path, test_loss=test_loss, test_loss1=test_loss1, test_loss2=test_loss2)

    def test_accuracy(self, model, test_loader):
        model.eval()
        acc = 0
        for i, (input_RSS, *_, target_noise_level) in enumerate(test_loader):
            input_RSS = input_RSS.to(self.device)
            *_, noise_pred = model(input_RSS)  # noise_pred: [batch_size, 3]
            noise_pred = noise_pred.detach().cpu().numpy()
            noise_pred = np.argmax(noise_pred, axis=1)
            labels = target_noise_level.detach().cpu().numpy()
            acc += np.sum((noise_pred == labels).astype(int))
        acc /= len(test_loader.dataset)
        return acc

    def test_model(self, noise_level='all'):
        model = MLVAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))
        model = model.to(self.device)
        model_path = self.path + 'model_' + str(self.args.n_epochs) + '.pth'
        model.load_state_dict(torch.load(model_path, map_location=self.device))
        model.eval()
        *_, test_loader, test_loader1, test_loader2 = self.load_data(noise_level)
        noise_class = 0  # noise_class = 0 stands for all noise classes
        if noise_level == 'low':
            noise_class = 1
        elif noise_level == 'mid':
            noise_class = 2
        elif noise_level == 'high':
            noise_class = 3

        # test model for testing set
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(
                test_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_pred = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader.dataset)
        ab_rmse /= len(test_loader.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader)
        print('====> Test for Testing Set:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))

        # test model for image 1
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader1):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_pred = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF1_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader1.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader1.dataset)
        ab_rmse /= len(test_loader1.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader1)
        print('====> Test for image 1:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))

        # test model for image 2
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader2):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_pred = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF2_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader2.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader2.dataset)
        ab_rmse /= len(test_loader2.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader2)
        print('====> Test for image 2:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))


# runner for Multi-task Gaussian Mixture VAE
class MGMVAE_runner():
    def __init__(self, args):
        self.args = args
        cuda_flag = not self.args.no_cuda and torch.cuda.is_available()
        self.device = torch.device("cuda" if cuda_flag else "cpu")
        torch.manual_seed(self.args.seed)
        random.seed(self.args.seed)
        self.path = self.args.run + 'MGMVAE/'
        if not os.path.exists(self.path):
            os.makedirs(self.path)

    def get_optimizer(self, parameters):
        if self.args.optimizer == 'Adam':
            return optim.Adam(parameters, lr=self.args.lr, weight_decay=self.args.weight_decay, betas=(0.9, 0.999), eps=1e-7)
        elif self.args.optimizer == 'RMSProp':
            return optim.RMSprop(parameters, lr=self.args.lr, weight_decay=self.args.weight_decay)
        elif self.args.optimizer == 'SGD':
            return optim.SGD(parameters, lr=self.args.lr, momentum=0.9)
        else:
            raise NotImplementedError('Optimizer {} not understood.'.format(self.args.optimizer))

    def load_data(self, noise_level='all'):
        dataset, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min = self.load_training_data()
        train_size = int(len(dataset) * 0.9)
        val_size = int(len(dataset) * 0.1)
        dataset_train, dataset_val = random_split(dataset, [train_size, val_size],
                                                  generator=torch.Generator().manual_seed(self.args.seed))
        dataset_test = self.load_testing_data(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        dataset_test1 = self.load_testing_img1(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        dataset_test2 = self.load_testing_img2(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        train_loader = DataLoader(
            dataset_train,
            batch_size=self.args.batch_size, shuffle=True, num_workers=2)
        val_loader = DataLoader(
            dataset_val,
            batch_size=self.args.batch_size, shuffle=True, num_workers=2)
        test_loader = DataLoader(
            dataset_test,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        test_loader1 = DataLoader(
            dataset_test1,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        test_loader2 = DataLoader(
            dataset_test2,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        print("Data Loaded!")
        return train_loader, val_loader, test_loader, test_loader1, test_loader2

    def load_training_data(self):
        # Load training data
        training_data = sio.loadmat('./data/training_data.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        train_input_RSS = training_data['RSS']
        RSS_max = np.amax(train_input_RSS, axis=0, keepdims=True)
        RSS_min = np.amin(train_input_RSS, axis=0, keepdims=True)
        train_input_RSS = (train_input_RSS - RSS_min) / (RSS_max - RSS_min)
        train_input_RSS = train_input_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        train_denoise_RSS = training_data['RSS_denoise']
        RSS_denoise_max = np.amax(train_denoise_RSS, axis=0, keepdims=True)
        RSS_denoise_min = np.amin(train_denoise_RSS, axis=0, keepdims=True)
        train_denoise_RSS = (train_denoise_RSS - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        train_denoise_RSS = train_denoise_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        train_slf_img = training_data['slf_img']
        train_slf_img = train_slf_img.reshape(-1, self.args.K0 * self.args.K1)
        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        train_ab = training_data['ab_norm']
        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        train_noise_level = training_data['sig_epsilon_class']
        train_noise_level = np.squeeze(train_noise_level)

        train_input_RSS = torch.from_numpy(train_input_RSS).float()
        train_denoise_RSS = torch.from_numpy(train_denoise_RSS).float()
        train_slf_img = torch.from_numpy(train_slf_img).float()
        train_ab = torch.from_numpy(train_ab).float()
        train_noise_level = torch.from_numpy(train_noise_level).long()

        Train_Dataset = TensorDataset(train_input_RSS, train_denoise_RSS, train_slf_img, train_ab, train_noise_level)
        return Train_Dataset, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min

    def load_testing_data(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data
        testing_data = sio.loadmat('./data/testing_data.mat')
        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_class = testing_data['sig_epsilon_class']
        test_noise_class = np.squeeze(test_noise_class)
        idx = (test_noise_class == 0) | (test_noise_class == 1) | (test_noise_class == 2)
        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            idx = (test_noise_class == 0)
        elif noise_level == 'mid':
            idx = (test_noise_class == 1)
        elif noise_level == 'high':
            idx = (test_noise_class == 2)
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))
        test_noise_class = test_noise_class[idx]
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS = testing_data['RSS'][idx]
        test_input_RSS = (test_input_RSS - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS = test_input_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS = testing_data['RSS_denoise'][idx]
        test_denoise_RSS = (test_denoise_RSS - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS = test_denoise_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img = testing_data['slf_img'][idx]
        test_slf_img = test_slf_img.reshape(-1, self.args.K0 * self.args.K1)

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab = testing_data['ab_norm'][idx]

        test_input_RSS = torch.from_numpy(test_input_RSS).float()
        test_denoise_RSS = torch.from_numpy(test_denoise_RSS).float()
        test_slf_img = torch.from_numpy(test_slf_img).float()
        test_ab = torch.from_numpy(test_ab).float()
        test_noise_class = torch.from_numpy(test_noise_class).long()

        Test_Dataset = TensorDataset(test_input_RSS, test_denoise_RSS, test_slf_img, test_ab, test_noise_class)
        return Test_Dataset

    def load_testing_img1(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data for image 1
        testing_data1 = sio.loadmat('./data/testing_data_image1.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS1 = testing_data1['RSS']
        test_input_RSS1 = (test_input_RSS1 - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS1 = test_input_RSS1.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS1 = testing_data1['RSS_denoise']
        test_denoise_RSS1 = (test_denoise_RSS1 - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS1 = test_denoise_RSS1.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img1 = testing_data1['slf_img']
        test_slf_img1 = test_slf_img1.reshape(1, self.args.K0 * self.args.K1)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        test_slf_img1 = np.tile(test_slf_img1, (test_input_RSS1.shape[0], 1))

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab1 = testing_data1['ab_norm']

        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_level1 = testing_data1['sig_epsilon_class']
        test_noise_level1 = np.squeeze(test_noise_level1)

        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            test_input_RSS1 = test_input_RSS1[:500]
            test_denoise_RSS1 = test_denoise_RSS1[:500]
            test_slf_img1 = test_slf_img1[:500]
            test_ab1 = test_ab1[:500]
            test_noise_level1 = test_noise_level1[:500]
        elif noise_level == 'mid':
            test_input_RSS1 = test_input_RSS1[500:1000]
            test_denoise_RSS1 = test_denoise_RSS1[500:1000]
            test_slf_img1 = test_slf_img1[500:1000]
            test_ab1 = test_ab1[500:1000]
            test_noise_level1 = test_noise_level1[500:1000]
        elif noise_level == 'high':
            test_input_RSS1 = test_input_RSS1[1000:]
            test_denoise_RSS1 = test_denoise_RSS1[1000:]
            test_slf_img1 = test_slf_img1[1000:]
            test_ab1 = test_ab1[1000:]
            test_noise_level1 = test_noise_level1[1000:]
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))

        test_input_RSS1 = torch.from_numpy(test_input_RSS1).float()
        test_denoise_RSS1 = torch.from_numpy(test_denoise_RSS1).float()
        test_slf_img1 = torch.from_numpy(test_slf_img1).float()
        test_ab1 = torch.from_numpy(test_ab1).float()
        test_noise_level1 = torch.from_numpy(test_noise_level1).long()

        Test_Dataset_1 = TensorDataset(test_input_RSS1, test_denoise_RSS1, test_slf_img1, test_ab1, test_noise_level1)
        return Test_Dataset_1

    def load_testing_img2(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data for image 2
        testing_data2 = sio.loadmat('./data/testing_data_image2.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS2 = testing_data2['RSS']
        test_input_RSS2 = (test_input_RSS2 - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS2 = test_input_RSS2.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS2 = testing_data2['RSS_denoise_norm']
        test_denoise_RSS2 = (test_denoise_RSS2 - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS2 = test_denoise_RSS2.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img2 = testing_data2['slf_img']
        test_slf_img2 = test_slf_img2.reshape(1, self.args.K0 * self.args.K1)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        test_slf_img2 = np.tile(test_slf_img2, (test_input_RSS2.shape[0], 1))

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab2 = testing_data2['ab_norm']

        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_level2 = testing_data2['sig_epsilon_class']
        test_noise_level2 = np.squeeze(test_noise_level2)

        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            test_input_RSS2 = test_input_RSS2[:500]
            test_denoise_RSS2 = test_denoise_RSS2[:500]
            test_slf_img2 = test_slf_img2[:500]
            test_ab2 = test_ab2[:500]
            test_noise_level2 = test_noise_level2[:500]
        elif noise_level == 'mid':
            test_input_RSS2 = test_input_RSS2[500:1000]
            test_denoise_RSS2 = test_denoise_RSS2[500:1000]
            test_slf_img2 = test_slf_img2[500:1000]
            test_ab2 = test_ab2[500:1000]
            test_noise_level2 = test_noise_level2[500:1000]
        elif noise_level == 'high':
            test_input_RSS2 = test_input_RSS2[1000:]
            test_denoise_RSS2 = test_denoise_RSS2[1000:]
            test_slf_img2 = test_slf_img2[1000:]
            test_ab2 = test_ab2[1000:]
            test_noise_level2 = test_noise_level2[1000:]
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))

        test_input_RSS2 = torch.from_numpy(test_input_RSS2).float()
        test_denoise_RSS2 = torch.from_numpy(test_denoise_RSS2).float()
        test_slf_img2 = torch.from_numpy(test_slf_img2).float()
        test_ab2 = torch.from_numpy(test_ab2).float()
        test_noise_level2 = torch.from_numpy(test_noise_level2).long()

        Test_Dataset_2 = TensorDataset(test_input_RSS2, test_denoise_RSS2, test_slf_img2, test_ab2, test_noise_level2)
        return Test_Dataset_2

    def train(self, model, train_loader, optimizer, epoch):
        model.train()
        train_loss = 0                  # total loss
        train_kld = 0                   # total kl-divergence
        train_loss_task0 = 0            # loss for reconstruct RSS
        train_loss_task1 = 0            # loss for SLF image estimation
        train_loss_task2 = 0            # loss for parameters estimation
        train_loss_task3 = 0            # loss for noise level prediction
        train_acc = 0                   # accuracy for noise level prediction
        train_kl_loss_w = 0             # kl-divergence for w
        train_kl_loss_h = 0             # kl-divergence for h
        train_kl_loss_c = 0             # kl-divergence for c
        batch_num = int(len(train_loader.dataset) / self.args.batch_size) + 1
        batch_10num = batch_num * 10
        for batch_idx,  (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(train_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            optimizer.zero_grad()
            KLD, RSS, SLF, ab, noise_level, kl_loss_w, kl_loss_h, kl_loss_c = model(input_RSS)
            loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum')/(self.args.N * self.args.P * self.args.P)
            loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum')/(self.args.K0 * self.args.K1)
            loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum')/(self.args.N + 1)
            loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
            loss_weight0 = (((epoch-1) % 20) * batch_num + batch_idx) / batch_10num
            loss_weight0 = np.minimum(loss_weight0, self.args.lambda0)
            loss = loss_weight0 * KLD + self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1\
                   + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3
            train_loss += loss.item()
            train_kld += KLD.item()
            train_kl_loss_w += kl_loss_w.item()
            train_kl_loss_h += kl_loss_h.item()
            train_kl_loss_c += kl_loss_c.item()
            train_loss_task0 += loss_task0.item()
            train_loss_task1 += loss_task1.item()
            train_loss_task2 += loss_task2.item()
            train_loss_task3 += loss_task3.item()
            _, noise_pred = torch.max(noise_level, 1)
            acc_task3 = (noise_pred == target_noise_level).sum()
            train_acc += acc_task3.item()
            loss /= input_RSS.size(0)
            loss.backward()
            optimizer.step()

            if batch_idx % self.args.log_interval == 0:
                print(
                    'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tKLD: {:.6f}\tLoss0: {:.6f}\tLoss1: {:.6f}\tLoss2: {:.6f}\tLoss3: {:.6f}\tAcc: {:.6f}'.format(
                        epoch, batch_idx * len(input_RSS), len(train_loader.dataset),
                               100. * batch_idx / len(train_loader),
                               loss.item(),
                               KLD.item() / len(input_RSS),
                               loss_task0.item() / len(input_RSS), loss_task1.item() / len(input_RSS),
                               loss_task2.item() / len(input_RSS), loss_task3.item() / len(input_RSS), acc_task3.item() / len(input_RSS)))
                print(
                    'Train Epoch: {} [{}/{} ({:.0f}%)]\tkld_w: {:.6f}\tkld_h: {:.6f}\tkld_c: {:.6f}'.format(
                        epoch, batch_idx * len(input_RSS), len(train_loader.dataset),
                               100. * batch_idx / len(train_loader),
                               kl_loss_w.item() / len(input_RSS),
                               kl_loss_h.item() / len(input_RSS),
                               kl_loss_c.item() / len(input_RSS)))

        train_loss /= len(train_loader.dataset)
        train_kld /= len(train_loader.dataset)
        train_kl_loss_w /= len(train_loader.dataset)
        train_kl_loss_h /= len(train_loader.dataset)
        train_kl_loss_c /= len(train_loader.dataset)
        train_loss_task0 /= len(train_loader.dataset)
        train_loss_task1 /= len(train_loader.dataset)
        train_loss_task2 /= len(train_loader.dataset)
        train_loss_task3 /= len(train_loader.dataset)
        train_acc /= len(train_loader.dataset)
        print('====> Epoch: {} Average loss: {:.6f}'.format(epoch, train_loss))
        print('====> KL divergence: {:.6f}'.format(train_kld))
        print('====> KLD_W: {:.6f}'.format(train_kl_loss_w))
        print('====> KLD_H: {:.6f}'.format(train_kl_loss_h))
        print('====> KLD_C: {:.6f}'.format(train_kl_loss_c))
        print('====> RSS loss: {:.6f}'.format(train_loss_task0))
        print('====> SLF loss: {:.6f}'.format(train_loss_task1))
        print('====> ab loss: {:.6f}'.format(train_loss_task2))
        print('====> noise level loss: {:.6f}'.format(train_loss_task3))
        print('====> noise level Accuracy: {:.4f}'.format(train_acc))
        return train_loss, train_kld, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3, train_acc, train_kl_loss_w, train_kl_loss_h, train_kl_loss_c

    def validate(self, model, val_loader, epoch):
        model.eval()
        val_loss = 0  # total loss
        val_rmse_slf = 0  # slf reconstruction loss
        val_mae_slf = 0  # slf reconstruction loss
        batch_num = int(len(val_loader.dataset) / self.args.batch_size) + 1
        batch_10num = batch_num * 10
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(val_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_level, *_ = model(input_RSS)
            loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum') / (
                        self.args.N * self.args.P * self.args.P)
            loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
            loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum') / (self.args.N + 1)
            loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
            loss_weight0 = (((epoch - 1) % 20) * batch_num + batch_idx) / batch_10num
            loss_weight0 = np.minimum(loss_weight0, self.args.lambda0)
            loss = 0 * KLD + self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1 \
                   + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3
            val_loss += loss.item()
            slf_mse = F.mse_loss(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
            slf_mae = F.l1_loss(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
            val_rmse_slf += slf_mse.item()
            val_mae_slf += slf_mae.item()
        val_loss /= len(val_loader.dataset)
        val_rmse_slf /= len(val_loader.dataset)
        val_rmse_slf = np.sqrt(val_rmse_slf)
        val_mae_slf /= len(val_loader.dataset)
        print('====> Validation set loss: {:.6f}'.format(val_loss))
        print('====> Validation set SLF rmse: {:.6f}'.format(val_rmse_slf))
        print('====> Validation set SLF mae: {:.6f}'.format(val_mae_slf))
        return val_loss, val_rmse_slf, val_mae_slf

    def test(self, model, test_loader, epoch):
        model.eval()
        test_loss = 0  # total loss
        test_kld = 0   # KL-divergence
        test_loss_task0 = 0  # loss for reconstruct RSS
        test_loss_task1 = 0  # loss for SLF image estimation
        test_loss_task2 = 0  # loss for parameters estimation
        test_loss_task3 = 0  # loss for noise level prediction
        test_kl_loss_w = 0  # kl-divergence for w
        test_kl_loss_h = 0  # kl-divergence for h
        test_kl_loss_c = 0  # kl-divergence for c
        batch_num = int(len(test_loader.dataset) / self.args.batch_size) + 1
        batch_10num = batch_num * 10
        with torch.no_grad():
            for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader):
                input_RSS = input_RSS.to(self.device)
                target_RSS = target_RSS.to(self.device)
                target_SLF = target_SLF.to(self.device)
                target_ab = target_ab.to(self.device)
                target_noise_level = target_noise_level.to(self.device)
                KLD, RSS, SLF, ab, noise_level, kl_loss_w, kl_loss_h, kl_loss_c = model(input_RSS)
                loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum') / (
                            self.args.N * self.args.P * self.args.P)
                loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
                loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum') / (self.args.N + 1)
                loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
                loss_weight0 = (((epoch - 1) % 20) * batch_num + batch_idx) / batch_10num
                loss_weight0 = np.minimum(loss_weight0, self.args.lambda0)
                loss = loss_weight0 * KLD + self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1 \
                       + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3

                test_loss += loss.item()
                test_kld += KLD.item()
                test_kl_loss_w += kl_loss_w.item()
                test_kl_loss_h += kl_loss_h.item()
                test_kl_loss_c += kl_loss_c.item()
                test_loss_task0 += loss_task0.item()
                test_loss_task1 += loss_task1.item()
                test_loss_task2 += loss_task2.item()
                test_loss_task3 += loss_task3.item()

        test_loss /= len(test_loader.dataset)
        test_kld /= len(test_loader.dataset)
        test_kl_loss_w /= len(test_loader.dataset)
        test_kl_loss_h /= len(test_loader.dataset)
        test_kl_loss_c /= len(test_loader.dataset)
        test_loss_task0 /= len(test_loader.dataset)
        test_loss_task1 /= len(test_loader.dataset)
        test_loss_task2 /= len(test_loader.dataset)
        test_loss_task3 /= len(test_loader.dataset)
        test_acc = self.test_accuracy(model, test_loader)
        print('====> Test set loss: {:.6f}'.format(test_loss))
        print('====> Test set KL divergence: {:.6f}'.format(test_kld))
        print('====> Test set KLD_W: {:.6f}'.format(test_kl_loss_w))
        print('====> Test set KLD_H: {:.6f}'.format(test_kl_loss_h))
        print('====> Test set KLD_C: {:.6f}'.format(test_kl_loss_c))
        print('====> Test set RSS loss: {:.6f}'.format(test_loss_task0))
        print('====> Test set SLF loss: {:.6f}'.format(test_loss_task1))
        print('====> Test set ab loss: {:.6f}'.format(test_loss_task2))
        print('====> Test set noise level loss: {:.6f}'.format(test_loss_task3))
        print('====> Test set noise level Accuracy: {:.4f}'.format(test_acc))
        return test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc, test_kl_loss_w, test_kl_loss_h, test_kl_loss_c

    def train_save(self):
        model = MGMVAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))

        def weights_init(m):
            if isinstance(m, torch.nn.Linear) or isinstance(m, torch.nn.Conv2d) or isinstance(m, torch.nn.ConvTranspose2d):
                torch.nn.init.xavier_uniform_(m.weight.data)
                if m.bias is not None:
                    torch.nn.init.zeros_(m.bias.data)

        model.apply(weights_init)
        model = model.to(self.device)
        summary(model, input_size=(self.args.batch_size, self.args.N, self.args.P, self.args.P))
        train_loader, val_loader, test_loader, test_loader1, test_loader2 = self.load_data()
        optimizer = self.get_optimizer(model.parameters())
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.25, patience=10, min_lr=1e-5)
        train_set = []
        val_set = []
        test_set = []      # for test set
        test_set1 = []      # for test image 1
        test_set2 = []      # for test image 2
        train_loss_path = self.path + 'train_loss_' + str(self.args.n_epochs) + '.npy'
        val_loss_path = self.path + 'val_loss_' + str(self.args.n_epochs) + '.npy'
        test_loss_path = self.path + 'test_loss_' + str(self.args.n_epochs) + '.npz'
        model_path = self.path + 'model_' + str(self.args.n_epochs) + '.pth'
        current_val_slf_loss = np.inf
        for epoch in range(1, self.args.n_epochs+1):
            print("Epoch %d learning rate：%f" % (epoch, optimizer.param_groups[0]['lr']))
            train_loss, train_kld, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3, train_acc, train_kl_loss_w, train_kl_loss_h, train_kl_loss_c = self.train(
                model,
                train_loader,
                optimizer,
                epoch)
            train_set.append(
                [train_loss, train_kld, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3,
                 train_acc, train_kl_loss_w, train_kl_loss_h, train_kl_loss_c])
            print('====> Validation Loss for Epoch {:d}'.format(epoch))
            val_loss, val_slf_rmse, val_slf_mae = self.validate(model, val_loader, epoch)
            val_set.append([val_loss, val_slf_rmse, val_slf_mae])
            scheduler.step(val_loss)
            print('====> Test for Epoch {:d}'.format(epoch))
            test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc, test_kl_loss_w, test_kl_loss_h, test_kl_loss_c = self.test(
                model,
                test_loader, epoch)
            test_set.append(
                [test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc,
                 test_kl_loss_w, test_kl_loss_h, test_kl_loss_c])
            print('====> Test Image1 for Epoch {:d}'.format(epoch))
            test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc, test_kl_loss_w, test_kl_loss_h, test_kl_loss_c = self.test(
                model,
                test_loader1, epoch)
            test_set1.append(
                [test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc,
                 test_kl_loss_w, test_kl_loss_h, test_kl_loss_c])
            print('====> Test Image2 for Epoch {:d}'.format(epoch))
            test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc, test_kl_loss_w, test_kl_loss_h, test_kl_loss_c = self.test(
                model,
                test_loader2, epoch)
            test_set2.append(
                [test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc,
                 test_kl_loss_w, test_kl_loss_h, test_kl_loss_c])
            if epoch % self.args.save_freq == 0:
                ckpt_path = self.path + '/model' + '_' + str(epoch) + '.pth'
                torch.save(model.state_dict(), ckpt_path)
                print('checkpoint{}.pth saved!'.format(epoch))
            if val_slf_mae < current_val_slf_loss:
                torch.save(model.state_dict(), model_path)
                current_val_slf_loss = val_slf_mae

        train_loss = np.asarray(train_set).reshape(-1, 10)
        val_loss = np.asarray(val_set).reshape(-1, 3)
        test_loss = np.asarray(test_set).reshape(-1, 10)
        test_loss1 = np.asarray(test_set1).reshape(-1, 10)
        test_loss2 = np.asarray(test_set2).reshape(-1, 10)
        np.save(train_loss_path, train_loss)
        np.save(val_loss_path, val_loss)
        np.savez(test_loss_path, test_loss=test_loss, test_loss1=test_loss1, test_loss2=test_loss2)

    def test_accuracy(self, model, test_loader):
        model.eval()
        acc = 0
        for i, (input_RSS, *_, target_noise_level) in enumerate(test_loader):
            input_RSS = input_RSS.to(self.device)
            KLD, RSS, SLF, ab, noise_pred, *_ = model(input_RSS)  # noise_pred: [batch_size, 3]
            noise_pred = noise_pred.detach().cpu().numpy()
            noise_pred = np.argmax(noise_pred, axis=1)
            labels = target_noise_level.detach().cpu().numpy()
            acc += np.sum((noise_pred == labels).astype(int))
        acc /= len(test_loader.dataset)
        return acc

    def test_model(self, noise_level='all'):
        model = MGMVAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))
        model = model.to(self.device)
        model_path = self.path + 'model_' + str(self.args.n_epochs) + '.pth'
        model.load_state_dict(torch.load(model_path, map_location=self.device))
        model.eval()
        *_, test_loader, test_loader1, test_loader2 = self.load_data(noise_level)
        noise_class = 0  # noise_class = 0 stands for all noise classes
        if noise_level == 'low':
            noise_class = 1
        elif noise_level == 'mid':
            noise_class = 2
        elif noise_level == 'high':
            noise_class = 3

        # test model for testing set
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(
                test_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_pred, *_ = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader.dataset)
        ab_rmse /= len(test_loader.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader)
        print('====> Test for Testing Set:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))

        # test model for image 1
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader1):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_pred, *_ = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF1_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader1.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader1.dataset)
        ab_rmse /= len(test_loader1.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader1)
        print('====> Test for image 1:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))

        # test model for image 2
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader2):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_pred, *_ = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF2_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader2.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader2.dataset)
        ab_rmse /= len(test_loader2.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader2)
        print('====> Test for image 2:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))


# runner for Multi-task Laplacian Mixture VAE
class MLMVAE_runner():
    def __init__(self, args):
        self.args = args
        cuda_flag = not self.args.no_cuda and torch.cuda.is_available()
        self.device = torch.device("cuda" if cuda_flag else "cpu")
        torch.manual_seed(self.args.seed)
        random.seed(self.args.seed)
        self.path = self.args.run + 'MLMVAE/'
        if not os.path.exists(self.path):
            os.makedirs(self.path)

    def get_optimizer(self, parameters):
        if self.args.optimizer == 'Adam':
            return optim.Adam(parameters, lr=self.args.lr, weight_decay=self.args.weight_decay, betas=(0.9, 0.999), eps=1e-7)
        elif self.args.optimizer == 'RMSProp':
            return optim.RMSprop(parameters, lr=self.args.lr, weight_decay=self.args.weight_decay)
        elif self.args.optimizer == 'SGD':
            return optim.SGD(parameters, lr=self.args.lr, momentum=0.9)
        else:
            raise NotImplementedError('Optimizer {} not understood.'.format(self.args.optimizer))

    def load_data(self, noise_level='all'):
        dataset, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min = self.load_training_data()
        train_size = int(len(dataset) * 0.9)
        val_size = int(len(dataset) * 0.1)
        dataset_train, dataset_val = random_split(dataset, [train_size, val_size],
                                                  generator=torch.Generator().manual_seed(self.args.seed))
        dataset_test = self.load_testing_data(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        dataset_test1 = self.load_testing_img1(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        dataset_test2 = self.load_testing_img2(RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level)
        train_loader = DataLoader(
            dataset_train,
            batch_size=self.args.batch_size, shuffle=True, num_workers=2)
        val_loader = DataLoader(
            dataset_val,
            batch_size=self.args.batch_size, shuffle=True, num_workers=2)
        test_loader = DataLoader(
            dataset_test,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        test_loader1 = DataLoader(
            dataset_test1,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        test_loader2 = DataLoader(
            dataset_test2,
            batch_size=self.args.batch_size, shuffle=False, num_workers=2)
        print("Data Loaded!")
        return train_loader, val_loader, test_loader, test_loader1, test_loader2

    def load_training_data(self):
        # Load training data
        training_data = sio.loadmat('./data/training_data.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        train_input_RSS = training_data['RSS']
        RSS_max = np.amax(train_input_RSS, axis=0, keepdims=True)
        RSS_min = np.amin(train_input_RSS, axis=0, keepdims=True)
        train_input_RSS = (train_input_RSS - RSS_min) / (RSS_max - RSS_min)
        train_input_RSS = train_input_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        train_denoise_RSS = training_data['RSS_denoise']
        RSS_denoise_max = np.amax(train_denoise_RSS, axis=0, keepdims=True)
        RSS_denoise_min = np.amin(train_denoise_RSS, axis=0, keepdims=True)
        train_denoise_RSS = (train_denoise_RSS - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        train_denoise_RSS = train_denoise_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        train_slf_img = training_data['slf_img']
        train_slf_img = train_slf_img.reshape(-1, self.args.K0 * self.args.K1)
        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        train_ab = training_data['ab_norm']
        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        train_noise_level = training_data['sig_epsilon_class']
        train_noise_level = np.squeeze(train_noise_level)

        train_input_RSS = torch.from_numpy(train_input_RSS).float()
        train_denoise_RSS = torch.from_numpy(train_denoise_RSS).float()
        train_slf_img = torch.from_numpy(train_slf_img).float()
        train_ab = torch.from_numpy(train_ab).float()
        train_noise_level = torch.from_numpy(train_noise_level).long()

        Train_Dataset = TensorDataset(train_input_RSS, train_denoise_RSS, train_slf_img, train_ab, train_noise_level)
        return Train_Dataset, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min

    def load_testing_data(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data
        testing_data = sio.loadmat('./data/testing_data.mat')
        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_class = testing_data['sig_epsilon_class']
        test_noise_class = np.squeeze(test_noise_class)
        idx = (test_noise_class == 0) | (test_noise_class == 1) | (test_noise_class == 2)
        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            idx = (test_noise_class == 0)
        elif noise_level == 'mid':
            idx = (test_noise_class == 1)
        elif noise_level == 'high':
            idx = (test_noise_class == 2)
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))
        test_noise_class = test_noise_class[idx]
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS = testing_data['RSS'][idx]
        test_input_RSS = (test_input_RSS - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS = test_input_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS = testing_data['RSS_denoise'][idx]
        test_denoise_RSS = (test_denoise_RSS - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS = test_denoise_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img = testing_data['slf_img'][idx]
        test_slf_img = test_slf_img.reshape(-1, self.args.K0 * self.args.K1)

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab = testing_data['ab_norm'][idx]

        test_input_RSS = torch.from_numpy(test_input_RSS).float()
        test_denoise_RSS = torch.from_numpy(test_denoise_RSS).float()
        test_slf_img = torch.from_numpy(test_slf_img).float()
        test_ab = torch.from_numpy(test_ab).float()
        test_noise_class = torch.from_numpy(test_noise_class).long()

        Test_Dataset = TensorDataset(test_input_RSS, test_denoise_RSS, test_slf_img, test_ab, test_noise_class)
        return Test_Dataset

    def load_testing_img1(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data for image 1
        testing_data1 = sio.loadmat('./data/testing_data_image1.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS1 = testing_data1['RSS']
        test_input_RSS1 = (test_input_RSS1 - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS1 = test_input_RSS1.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS1 = testing_data1['RSS_denoise']
        test_denoise_RSS1 = (test_denoise_RSS1 - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS1 = test_denoise_RSS1.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img1 = testing_data1['slf_img']
        test_slf_img1 = test_slf_img1.reshape(1, self.args.K0 * self.args.K1)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        test_slf_img1 = np.tile(test_slf_img1, (test_input_RSS1.shape[0], 1))

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab1 = testing_data1['ab_norm']

        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_level1 = testing_data1['sig_epsilon_class']
        test_noise_level1 = np.squeeze(test_noise_level1)

        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            test_input_RSS1 = test_input_RSS1[:500]
            test_denoise_RSS1 = test_denoise_RSS1[:500]
            test_slf_img1 = test_slf_img1[:500]
            test_ab1 = test_ab1[:500]
            test_noise_level1 = test_noise_level1[:500]
        elif noise_level == 'mid':
            test_input_RSS1 = test_input_RSS1[500:1000]
            test_denoise_RSS1 = test_denoise_RSS1[500:1000]
            test_slf_img1 = test_slf_img1[500:1000]
            test_ab1 = test_ab1[500:1000]
            test_noise_level1 = test_noise_level1[500:1000]
        elif noise_level == 'high':
            test_input_RSS1 = test_input_RSS1[1000:]
            test_denoise_RSS1 = test_denoise_RSS1[1000:]
            test_slf_img1 = test_slf_img1[1000:]
            test_ab1 = test_ab1[1000:]
            test_noise_level1 = test_noise_level1[1000:]
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))

        test_input_RSS1 = torch.from_numpy(test_input_RSS1).float()
        test_denoise_RSS1 = torch.from_numpy(test_denoise_RSS1).float()
        test_slf_img1 = torch.from_numpy(test_slf_img1).float()
        test_ab1 = torch.from_numpy(test_ab1).float()
        test_noise_level1 = torch.from_numpy(test_noise_level1).long()

        Test_Dataset_1 = TensorDataset(test_input_RSS1, test_denoise_RSS1, test_slf_img1, test_ab1, test_noise_level1)
        return Test_Dataset_1

    def load_testing_img2(self, RSS_max, RSS_min, RSS_denoise_max, RSS_denoise_min, noise_level='all'):
        # Load testing data for image 2
        testing_data2 = sio.loadmat('./data/testing_data_image2.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_input_RSS2 = testing_data2['RSS']
        test_input_RSS2 = (test_input_RSS2 - RSS_min) / (RSS_max - RSS_min)
        test_input_RSS2 = test_input_RSS2.reshape(-1, self.args.N, self.args.P, self.args.P)

        # normalized clean RSS measurement input (as target for decoder) (shape: [num_sample, NxPxP])
        test_denoise_RSS2 = testing_data2['RSS_denoise_norm']
        test_denoise_RSS2 = (test_denoise_RSS2 - RSS_denoise_min) / (RSS_denoise_max - RSS_denoise_min)
        test_denoise_RSS2 = test_denoise_RSS2.reshape(-1, self.args.N, self.args.P, self.args.P)

        # SLF image (shape: [1, K[0]*K[1]])
        test_slf_img2 = testing_data2['slf_img']
        test_slf_img2 = test_slf_img2.reshape(1, self.args.K0 * self.args.K1)
        # SLF image (shape: [num_sample, K[0]*K[1]])
        test_slf_img2 = np.tile(test_slf_img2, (test_input_RSS2.shape[0], 1))

        # b: bias. alpha: path loss exponent. ab_norm: [b_norm, alpha_norm] (shape: [num_sample, N+1])
        test_ab2 = testing_data2['ab_norm']

        # noise level classes (shape: [num_sample, 1]) (classes: 0, 1, 2)
        test_noise_level2 = testing_data2['sig_epsilon_class']
        test_noise_level2 = np.squeeze(test_noise_level2)

        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            test_input_RSS2 = test_input_RSS2[:500]
            test_denoise_RSS2 = test_denoise_RSS2[:500]
            test_slf_img2 = test_slf_img2[:500]
            test_ab2 = test_ab2[:500]
            test_noise_level2 = test_noise_level2[:500]
        elif noise_level == 'mid':
            test_input_RSS2 = test_input_RSS2[500:1000]
            test_denoise_RSS2 = test_denoise_RSS2[500:1000]
            test_slf_img2 = test_slf_img2[500:1000]
            test_ab2 = test_ab2[500:1000]
            test_noise_level2 = test_noise_level2[500:1000]
        elif noise_level == 'high':
            test_input_RSS2 = test_input_RSS2[1000:]
            test_denoise_RSS2 = test_denoise_RSS2[1000:]
            test_slf_img2 = test_slf_img2[1000:]
            test_ab2 = test_ab2[1000:]
            test_noise_level2 = test_noise_level2[1000:]
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))

        test_input_RSS2 = torch.from_numpy(test_input_RSS2).float()
        test_denoise_RSS2 = torch.from_numpy(test_denoise_RSS2).float()
        test_slf_img2 = torch.from_numpy(test_slf_img2).float()
        test_ab2 = torch.from_numpy(test_ab2).float()
        test_noise_level2 = torch.from_numpy(test_noise_level2).long()

        Test_Dataset_2 = TensorDataset(test_input_RSS2, test_denoise_RSS2, test_slf_img2, test_ab2, test_noise_level2)
        return Test_Dataset_2

    def train(self, model, train_loader, optimizer, epoch):
        model.train()
        train_loss = 0                  # total loss
        train_kld = 0                   # total kl-divergence
        train_loss_task0 = 0            # loss for reconstruct RSS
        train_loss_task1 = 0            # loss for SLF image estimation
        train_loss_task2 = 0            # loss for parameters estimation
        train_loss_task3 = 0            # loss for noise level prediction
        train_acc = 0                   # accuracy for noise level prediction
        train_kl_loss_w = 0             # kl-divergence for w
        train_kl_loss_h = 0             # kl-divergence for h
        train_kl_loss_c = 0             # kl-divergence for c
        batch_num = int(len(train_loader.dataset) / self.args.batch_size) + 1
        batch_10num = batch_num * 10
        for batch_idx,  (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(train_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            optimizer.zero_grad()
            KLD, RSS, SLF, ab, noise_level, kl_loss_w, kl_loss_h, kl_loss_c = model(input_RSS)
            loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum')/(self.args.N * self.args.P * self.args.P)
            loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum')/(self.args.K0 * self.args.K1)
            loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum')/(self.args.N + 1)
            loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
            loss_weight0 = (((epoch-1) % 20) * batch_num + batch_idx) / batch_10num
            loss_weight0 = np.minimum(loss_weight0, self.args.lambda0)
            loss = loss_weight0 * KLD + self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1\
                   + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3
            train_loss += loss.item()
            train_kld += KLD.item()
            train_kl_loss_w += kl_loss_w.item()
            train_kl_loss_h += kl_loss_h.item()
            train_kl_loss_c += kl_loss_c.item()
            train_loss_task0 += loss_task0.item()
            train_loss_task1 += loss_task1.item()
            train_loss_task2 += loss_task2.item()
            train_loss_task3 += loss_task3.item()
            _, noise_pred = torch.max(noise_level, 1)
            acc_task3 = (noise_pred == target_noise_level).sum()
            train_acc += acc_task3.item()
            loss /= input_RSS.size(0)
            loss.backward()
            optimizer.step()

            if batch_idx % self.args.log_interval == 0:
                print(
                    'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tKLD: {:.6f}\tLoss0: {:.6f}\tLoss1: {:.6f}\tLoss2: {:.6f}\tLoss3: {:.6f}\tAcc: {:.6f}'.format(
                        epoch, batch_idx * len(input_RSS), len(train_loader.dataset),
                               100. * batch_idx / len(train_loader),
                               loss.item(),
                               KLD.item() / len(input_RSS),
                               loss_task0.item() / len(input_RSS), loss_task1.item() / len(input_RSS),
                               loss_task2.item() / len(input_RSS), loss_task3.item() / len(input_RSS), acc_task3.item() / len(input_RSS)))
                print(
                    'Train Epoch: {} [{}/{} ({:.0f}%)]\tkld_w: {:.6f}\tkld_h: {:.6f}\tkld_c: {:.6f}'.format(
                        epoch, batch_idx * len(input_RSS), len(train_loader.dataset),
                               100. * batch_idx / len(train_loader),
                               kl_loss_w.item() / len(input_RSS),
                               kl_loss_h.item() / len(input_RSS),
                               kl_loss_c.item() / len(input_RSS)))

        train_loss /= len(train_loader.dataset)
        train_kld /= len(train_loader.dataset)
        train_kl_loss_w /= len(train_loader.dataset)
        train_kl_loss_h /= len(train_loader.dataset)
        train_kl_loss_c /= len(train_loader.dataset)
        train_loss_task0 /= len(train_loader.dataset)
        train_loss_task1 /= len(train_loader.dataset)
        train_loss_task2 /= len(train_loader.dataset)
        train_loss_task3 /= len(train_loader.dataset)
        train_acc /= len(train_loader.dataset)
        print('====> Epoch: {} Average loss: {:.6f}'.format(epoch, train_loss))
        print('====> KL divergence: {:.6f}'.format(train_kld))
        print('====> KLD_W: {:.6f}'.format(train_kl_loss_w))
        print('====> KLD_H: {:.6f}'.format(train_kl_loss_h))
        print('====> KLD_C: {:.6f}'.format(train_kl_loss_c))
        print('====> RSS loss: {:.6f}'.format(train_loss_task0))
        print('====> SLF loss: {:.6f}'.format(train_loss_task1))
        print('====> ab loss: {:.6f}'.format(train_loss_task2))
        print('====> noise level loss: {:.6f}'.format(train_loss_task3))
        print('====> noise level Accuracy: {:.4f}'.format(train_acc))
        return train_loss, train_kld, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3, train_acc, train_kl_loss_w, train_kl_loss_h, train_kl_loss_c

    def validate(self, model, val_loader, epoch):
        model.eval()
        val_loss = 0  # total loss
        val_rmse_slf = 0  # slf reconstruction loss
        val_mae_slf = 0  # slf reconstruction loss
        batch_num = int(len(val_loader.dataset) / self.args.batch_size) + 1
        batch_10num = batch_num * 10
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(val_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_level, *_ = model(input_RSS)
            loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum') / (
                    self.args.N * self.args.P * self.args.P)
            loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
            loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum') / (self.args.N + 1)
            loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
            loss_weight0 = (((epoch - 1) % 20) * batch_num + batch_idx) / batch_10num
            loss_weight0 = np.minimum(loss_weight0, self.args.lambda0)
            loss = 0 * KLD + self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1 \
                   + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3
            val_loss += loss.item()
            slf_mse = F.mse_loss(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
            slf_mae = F.l1_loss(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
            val_rmse_slf += slf_mse.item()
            val_mae_slf += slf_mae.item()
        val_loss /= len(val_loader.dataset)
        val_rmse_slf /= len(val_loader.dataset)
        val_rmse_slf = np.sqrt(val_rmse_slf)
        val_mae_slf /= len(val_loader.dataset)
        print('====> Validation set loss: {:.6f}'.format(val_loss))
        print('====> Validation set SLF rmse: {:.6f}'.format(val_rmse_slf))
        print('====> Validation set SLF mae: {:.6f}'.format(val_mae_slf))
        return val_loss, val_rmse_slf, val_mae_slf

    def test(self, model, test_loader, epoch):
        model.eval()
        test_loss = 0  # total loss
        test_kld = 0   # KL-divergence
        test_loss_task0 = 0  # loss for reconstruct RSS
        test_loss_task1 = 0  # loss for SLF image estimation
        test_loss_task2 = 0  # loss for parameters estimation
        test_loss_task3 = 0  # loss for noise level prediction
        test_kl_loss_w = 0  # kl-divergence for w
        test_kl_loss_h = 0  # kl-divergence for h
        test_kl_loss_c = 0  # kl-divergence for c
        batch_num = int(len(test_loader.dataset) / self.args.batch_size) + 1
        batch_10num = batch_num * 10
        with torch.no_grad():
            for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader):
                input_RSS = input_RSS.to(self.device)
                target_RSS = target_RSS.to(self.device)
                target_SLF = target_SLF.to(self.device)
                target_ab = target_ab.to(self.device)
                target_noise_level = target_noise_level.to(self.device)
                KLD, RSS, SLF, ab, noise_level, kl_loss_w, kl_loss_h, kl_loss_c = model(input_RSS)
                loss_task0 = F.binary_cross_entropy(RSS, target_RSS, reduction='sum') / (
                            self.args.N * self.args.P * self.args.P)
                loss_task1 = F.binary_cross_entropy(SLF, target_SLF, reduction='sum') / (self.args.K0 * self.args.K1)
                loss_task2 = F.binary_cross_entropy(ab, target_ab, reduction='sum') / (self.args.N + 1)
                loss_task3 = F.cross_entropy(noise_level, target_noise_level, reduction='sum')
                loss_weight0 = (((epoch - 1) % 20) * batch_num + batch_idx) / batch_10num
                loss_weight0 = np.minimum(loss_weight0, self.args.lambda0)
                loss = loss_weight0 * KLD + self.args.lambda1 * loss_task0 + self.args.lambda2 * loss_task1 \
                       + self.args.lambda3 * loss_task2 + self.args.lambda4 * loss_task3

                test_loss += loss.item()
                test_kld += KLD.item()
                test_kl_loss_w += kl_loss_w.item()
                test_kl_loss_h += kl_loss_h.item()
                test_kl_loss_c += kl_loss_c.item()
                test_loss_task0 += loss_task0.item()
                test_loss_task1 += loss_task1.item()
                test_loss_task2 += loss_task2.item()
                test_loss_task3 += loss_task3.item()

        test_loss /= len(test_loader.dataset)
        test_kld /= len(test_loader.dataset)
        test_kl_loss_w /= len(test_loader.dataset)
        test_kl_loss_h /= len(test_loader.dataset)
        test_kl_loss_c /= len(test_loader.dataset)
        test_loss_task0 /= len(test_loader.dataset)
        test_loss_task1 /= len(test_loader.dataset)
        test_loss_task2 /= len(test_loader.dataset)
        test_loss_task3 /= len(test_loader.dataset)
        test_acc = self.test_accuracy(model, test_loader)
        print('====> Test set loss: {:.6f}'.format(test_loss))
        print('====> Test set KL divergence: {:.6f}'.format(test_kld))
        print('====> Test set KLD_W: {:.6f}'.format(test_kl_loss_w))
        print('====> Test set KLD_H: {:.6f}'.format(test_kl_loss_h))
        print('====> Test set KLD_C: {:.6f}'.format(test_kl_loss_c))
        print('====> Test set RSS loss: {:.6f}'.format(test_loss_task0))
        print('====> Test set SLF loss: {:.6f}'.format(test_loss_task1))
        print('====> Test set ab loss: {:.6f}'.format(test_loss_task2))
        print('====> Test set noise level loss: {:.6f}'.format(test_loss_task3))
        print('====> Test set noise level Accuracy: {:.4f}'.format(test_acc))
        return test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc, test_kl_loss_w, test_kl_loss_h, test_kl_loss_c

    def train_save(self):
        model = MLMVAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))

        def weights_init(m):
            if isinstance(m, torch.nn.Linear) or isinstance(m, torch.nn.Conv2d) or isinstance(m,
                                                                                              torch.nn.ConvTranspose2d):
                torch.nn.init.xavier_uniform_(m.weight.data)
                if m.bias is not None:
                    torch.nn.init.zeros_(m.bias.data)

        model.apply(weights_init)
        model = model.to(self.device)
        summary(model, input_size=(self.args.batch_size, self.args.N, self.args.P, self.args.P))
        train_loader, val_loader, test_loader, test_loader1, test_loader2 = self.load_data()
        optimizer = self.get_optimizer(model.parameters())
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.25, patience=10, min_lr=1e-5)
        train_set = []
        val_set = []
        test_set = []  # for test set
        test_set1 = []  # for test image 1
        test_set2 = []  # for test image 2
        train_loss_path = self.path + 'train_loss_' + str(self.args.n_epochs) + '.npy'
        val_loss_path = self.path + 'val_loss_' + str(self.args.n_epochs) + '.npy'
        test_loss_path = self.path + 'test_loss_' + str(self.args.n_epochs) + '.npz'
        model_path = self.path + 'model_' + str(self.args.n_epochs) + '.pth'
        current_val_slf_loss = np.inf
        for epoch in range(1, self.args.n_epochs + 1):
            print("Epoch %d learning rate：%f" % (epoch, optimizer.param_groups[0]['lr']))
            train_loss, train_kld, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3, train_acc, train_kl_loss_w, train_kl_loss_h, train_kl_loss_c = self.train(
                model,
                train_loader,
                optimizer,
                epoch)
            train_set.append(
                [train_loss, train_kld, train_loss_task0, train_loss_task1, train_loss_task2, train_loss_task3,
                 train_acc, train_kl_loss_w, train_kl_loss_h, train_kl_loss_c])
            print('====> Validation Loss for Epoch {:d}'.format(epoch))
            val_loss, val_slf_rmse, val_slf_mae = self.validate(model, val_loader, epoch)
            val_set.append([val_loss, val_slf_rmse, val_slf_mae])
            scheduler.step(val_loss)
            print('====> Test for Epoch {:d}'.format(epoch))
            test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc, test_kl_loss_w, test_kl_loss_h, test_kl_loss_c = self.test(
                model,
                test_loader, epoch)
            test_set.append(
                [test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc,
                 test_kl_loss_w, test_kl_loss_h, test_kl_loss_c])
            print('====> Test Image1 for Epoch {:d}'.format(epoch))
            test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc, test_kl_loss_w, test_kl_loss_h, test_kl_loss_c = self.test(
                model,
                test_loader1, epoch)
            test_set1.append(
                [test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc,
                 test_kl_loss_w, test_kl_loss_h, test_kl_loss_c])
            print('====> Test Image2 for Epoch {:d}'.format(epoch))
            test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc, test_kl_loss_w, test_kl_loss_h, test_kl_loss_c = self.test(
                model,
                test_loader2, epoch)
            test_set2.append(
                [test_loss, test_kld, test_loss_task0, test_loss_task1, test_loss_task2, test_loss_task3, test_acc,
                 test_kl_loss_w, test_kl_loss_h, test_kl_loss_c])
            if epoch % self.args.save_freq == 0:
                ckpt_path = self.path + '/model' + '_' + str(epoch) + '.pth'
                torch.save(model.state_dict(), ckpt_path)
                print('checkpoint{}.pth saved!'.format(epoch))
            if val_slf_mae < current_val_slf_loss:
                torch.save(model.state_dict(), model_path)
                current_val_slf_loss = val_slf_mae

        train_loss = np.asarray(train_set).reshape(-1, 10)
        val_loss = np.asarray(val_set).reshape(-1, 3)
        test_loss = np.asarray(test_set).reshape(-1, 10)
        test_loss1 = np.asarray(test_set1).reshape(-1, 10)
        test_loss2 = np.asarray(test_set2).reshape(-1, 10)
        np.save(train_loss_path, train_loss)
        np.save(val_loss_path, val_loss)
        np.savez(test_loss_path, test_loss=test_loss, test_loss1=test_loss1, test_loss2=test_loss2)

    def test_accuracy(self, model, test_loader):
        model.eval()
        acc = 0
        for i, (input_RSS, *_, target_noise_level) in enumerate(test_loader):
            input_RSS = input_RSS.to(self.device)
            KLD, RSS, SLF, ab, noise_pred, *_ = model(input_RSS)  # noise_pred: [batch_size, 3]
            noise_pred = noise_pred.detach().cpu().numpy()
            noise_pred = np.argmax(noise_pred, axis=1)
            labels = target_noise_level.detach().cpu().numpy()
            acc += np.sum((noise_pred == labels).astype(int))
        acc /= len(test_loader.dataset)
        return acc

    def test_model(self, noise_level='all'):
        model = MLMVAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))
        model = model.to(self.device)
        model_path = self.path + 'model_' + str(self.args.n_epochs) + '.pth'
        model.load_state_dict(torch.load(model_path, map_location=self.device))
        model.eval()
        *_, test_loader, test_loader1, test_loader2 = self.load_data(noise_level)
        noise_class = 0  # noise_class = 0 stands for all noise classes
        if noise_level == 'low':
            noise_class = 1
        elif noise_level == 'mid':
            noise_class = 2
        elif noise_level == 'high':
            noise_class = 3

        # test model for testing set
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(
                test_loader):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_pred, *_ = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader.dataset)
        ab_rmse /= len(test_loader.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader)
        print('====> Test for Testing Set:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))

        # test model for image 1
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader1):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_pred, *_ = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF1_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader1.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader1.dataset)
        ab_rmse /= len(test_loader1.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader1)
        print('====> Test for image 1:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))

        # test model for image 2
        slf_rmse = 0
        slf_mae = 0
        ab_rmse = 0
        for batch_idx, (input_RSS, target_RSS, target_SLF, target_ab, target_noise_level) in enumerate(test_loader2):
            input_RSS = input_RSS.to(self.device)
            target_RSS = target_RSS.to(self.device)
            target_SLF = target_SLF.to(self.device)
            target_ab = target_ab.to(self.device)
            target_noise_level = target_noise_level.to(self.device)
            KLD, RSS, SLF, ab, noise_pred, *_ = model(input_RSS)
            # plot SLF reconstruction image for image1
            n = min(target_RSS.size(0), 8)
            comparison = torch.cat([target_SLF.view(target_RSS.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    SLF.view(SLF.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconstruction_SLF2_noise' + str(noise_class) + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse += F.mse_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae += F.l1_loss(SLF, target_SLF, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test ab parameters RMSE
            ab_rmse += F.mse_loss(ab, target_ab, reduction='sum').item() / (self.args.N + 1)

        slf_rmse /= len(test_loader2.dataset)
        slf_rmse = np.sqrt(slf_rmse)
        slf_mae /= len(test_loader2.dataset)
        ab_rmse /= len(test_loader2.dataset)
        ab_rmse = np.sqrt(ab_rmse)
        noise_acc = self.test_accuracy(model, test_loader2)
        print('====> Test for image 2:')
        print('====> Noise_level: ' + noise_level)
        print('====> Test set SLF rmse: {:.6f}'.format(slf_rmse))
        print('====> Test set SLF mae: {:.6f}'.format(slf_mae))
        print('====> Test set ab rmse: {:.6f}'.format(ab_rmse))
        print('====> Test set noise level Accuracy: {:.4f}'.format(noise_acc))