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 import SLF_Estimator, SLF_VAE
import sys
sys.path.append('..')
import random
import scipy.io as sio
import cv2


# runner for CNN-attention estimator
class SLF_Estimator_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 + 'SLF_Estimator/'
        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))
        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'):
        # Load training data
        training_data = sio.loadmat('./data/training_data.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        train_RSS = training_data['RSS']
        RSS_min = np.amin(train_RSS, axis=0, keepdims=True)
        RSS_max = np.amax(train_RSS, axis=0, keepdims=True)
        train_RSS = (train_RSS - RSS_min) / (RSS_max - RSS_min)
        train_RSS = train_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # SLF image (shape: [num_sample, K0, K1])
        train_slf_img = training_data['slf_img']
        train_slf_img = train_slf_img.reshape(-1, 1, self.args.K0, self.args.K1)

        train_RSS = torch.from_numpy(train_RSS).float()
        train_slf_img = torch.from_numpy(train_slf_img).float()

        dataset = TensorDataset(train_RSS, train_slf_img)
        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))

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

        # SLF image ENR (shape: [num_sample, K0, K1])
        test_slf_img_ENR = testing_data['slf_img_ENR'][idx]
        test_slf_img_ENR = test_slf_img_ENR.reshape(-1, 1, self.args.K0, self.args.K1)

        test_RSS = torch.from_numpy(test_RSS).float()
        test_slf_img = torch.from_numpy(test_slf_img).float()
        test_slf_img_ENR = torch.from_numpy(test_slf_img_ENR).float()

        dataset_test = TensorDataset(test_RSS, test_slf_img_ENR, test_slf_img)

        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)
        print("Data Loaded!")
        return train_loader, val_loader, test_loader

    def train(self, model, train_loader, optimizer, epoch):
        model.train()
        train_recon_loss = 0                  # total loss
        train_rmse = 0                  # rmse loss
        for batch_idx,  (input_rss, target_slf) in enumerate(train_loader):
            input_rss = input_rss.to(self.device)
            target_slf = target_slf.to(self.device)
            optimizer.zero_grad()
            slf_recon = model(input_rss)
            loss = F.binary_cross_entropy(slf_recon, target_slf, reduction='sum')/(self.args.K0*self.args.K1)
            train_recon_loss += loss.item()

            mse = F.mse_loss(slf_recon, target_slf, reduction='sum') / (self.args.K0 * self.args.K1)
            train_rmse += mse.item()

            loss /= input_rss.size(0)
            loss.backward()
            optimizer.step()
            if batch_idx % self.args.log_interval == 0:
                print(
                    'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                        epoch, batch_idx * len(input_rss), len(train_loader.dataset),
                               100. * batch_idx / len(train_loader),
                               loss.item()))

        train_recon_loss /= len(train_loader.dataset)
        train_rmse /= len(train_loader.dataset)
        train_rmse = np.sqrt(train_rmse)
        print('====> Epoch: {} Average loss: {:.6f}'.format(epoch, train_recon_loss))
        print('====> RMSE: {:.6f}'.format(train_rmse))
        return train_recon_loss, train_rmse

    def validate(self, model, val_loader):
        model.eval()
        val_recon_loss = 0           # recon loss
        val_rmse = 0           # rmse
        for batch_idx,  (input_rss, target_slf) in enumerate(val_loader):
            input_rss = input_rss.to(self.device)
            target_slf = target_slf.to(self.device)
            slf_recon = model(input_rss)
            recon_loss = F.binary_cross_entropy(slf_recon, target_slf, reduction='sum')/(self.args.K0*self.args.K1)
            val_recon_loss += recon_loss.item()
            mse = F.mse_loss(slf_recon, target_slf, reduction='sum')/(self.args.K0*self.args.K1)
            val_rmse += mse.item()
        val_recon_loss /= len(val_loader.dataset)
        val_rmse /= len(val_loader.dataset)
        val_rmse = np.sqrt(val_rmse)
        print('====> Loss: {:.6f}'.format(val_recon_loss))
        print('====> RMSE: {:.6f}'.format(val_rmse))
        return val_recon_loss, val_rmse

    def test(self, model, test_loader):
        model.eval()
        test_recon_loss = 0  # recon loss
        test_rmse1 = 0  # rmse for ENR
        test_rmse2 = 0  # rmse for our model
        for batch_idx, (input_rss, slf_enr, target_slf) in enumerate(test_loader):
            input_rss = input_rss.to(self.device)
            slf_enr = slf_enr.to(self.device)
            target_slf = target_slf.to(self.device)
            slf_recon = model(input_rss)
            recon_loss = F.binary_cross_entropy(slf_recon, target_slf, reduction='sum')/(self.args.K0*self.args.K1)
            test_recon_loss += recon_loss.item()
            mse1 = F.mse_loss(slf_enr, target_slf, reduction='sum') / (self.args.K0 * self.args.K1)
            mse2 = F.mse_loss(slf_recon, target_slf, reduction='sum') / (self.args.K0 * self.args.K1)
            test_rmse1 += mse1.item()
            test_rmse2 += mse2.item()
        test_recon_loss /= len(test_loader.dataset)
        test_rmse1 /= len(test_loader.dataset)
        test_rmse1 = np.sqrt(test_rmse1)
        test_rmse2 /= len(test_loader.dataset)
        test_rmse2 = np.sqrt(test_rmse2)
        print('====> Loss: {:.6f}'.format(test_recon_loss))
        print('====> RMSE for ENR: {:.6f}'.format(test_rmse1))
        print('====> RMSE for our model: {:.6f}'.format(test_rmse2))
        return test_recon_loss, test_rmse1, test_rmse2

    def train_save(self):
        model = SLF_Estimator(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))
        model = model.to(self.device)
        train_loader, val_loader, test_loader = self.load_data()
        optimizer = self.get_optimizer(model.parameters())
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.2, patience=8, min_lr=1e-5)
        train_set = []
        val_set = []
        test_set = []
        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) + '.npy'
        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_recon_loss, train_rmse = self.train(model, train_loader, optimizer, epoch)
            train_set.append([train_recon_loss, train_rmse])
            print('====> Validation Loss for Epoch {:d}'.format(epoch))
            val_recon_loss, val_rmse = self.validate(model, val_loader)
            val_set.append([val_recon_loss, val_rmse])
            scheduler.step(val_recon_loss)
            print('====> Test Loss for Epoch {:d}'.format(epoch))
            test_recon_loss, test_rmse1, test_rmse2 = self.test(model, test_loader)
            test_set.append([test_recon_loss, test_rmse1, test_rmse2])
            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_rmse < current_val_slf_loss:
                torch.save(model.state_dict(), model_path)
                current_val_slf_loss = val_rmse

        train_loss = np.asarray(train_set).reshape(-1, 2)
        val_loss = np.asarray(val_set).reshape(-1, 2)
        test_loss = np.asarray(test_set).reshape(-1, 3)
        np.save(train_loss_path, train_loss)
        np.save(val_loss_path, val_loss)
        np.save(test_loss_path, test_loss)

    def test_model(self, noise_level='all'):
        model = SLF_Estimator(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(noise_level)
        # test model
        slf_rmse1 = 0        # rmse for stage 1
        slf_mae1 = 0         # mae for stage 1
        slf_rmse2 = 0        # rmse for stage 2
        slf_mae2 = 0         # mae for stage 2
        for batch_idx, (input_rss, slf_enr, target_slf) in enumerate(test_loader):
            input_rss = input_rss.to(self.device)
            slf_enr = slf_enr.to(self.device)
            target_slf = target_slf.to(self.device)
            slf_recon = model(input_rss)
            # plot SLF reconstruction image for image1
            n = min(input_rss.size(0), 8)
            comparison = torch.cat([target_slf.view(target_slf.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    slf_enr.view(slf_enr.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    slf_recon.view(slf_recon.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconSLF_' + noise_level + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse1 += F.mse_loss(slf_enr, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            slf_rmse2 += F.mse_loss(slf_recon, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae1 += F.l1_loss(slf_enr, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            slf_mae2 += F.l1_loss(slf_recon, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)

        slf_rmse1 /= len(test_loader.dataset)
        slf_rmse1 = np.sqrt(slf_rmse1)
        slf_rmse2 /= len(test_loader.dataset)
        slf_rmse2 = np.sqrt(slf_rmse2)
        slf_mae1 /= len(test_loader.dataset)
        slf_mae2 /= len(test_loader.dataset)
        print('====> Test Model:')
        print('====> Noise Level: ' + noise_level)
        print('====> ENR SLF rmse: {:.6f}'.format(slf_rmse1))
        print('====> Recon SLF rmse: {:.6f}'.format(slf_rmse2))
        print('====> ENR SLF mae: {:.6f}'.format(slf_mae1))
        print('====> Recon SLF mae: {:.6f}'.format(slf_mae2))

    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_RSS = training_data['RSS']
        RSS_min = np.amin(train_RSS, axis=0, keepdims=True)
        RSS_max = np.amax(train_RSS, axis=0, keepdims=True)

        # Load testing data for statistical results simulation (ROC and mAP)
        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)
        noise_class = testing_data['sig_epsilon_class']
        noise_class = np.squeeze(noise_class)
        idx = (noise_class == 0) | (noise_class == 1) | (noise_class == 2)
        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            idx = (noise_class == 0)
        elif noise_level == 'mid':
            idx = (noise_class == 1)
        elif noise_level == 'high':
            idx = (noise_class == 2)
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_RSS = testing_data['RSS'][idx]
        test_RSS = (test_RSS - RSS_min) / (RSS_max - RSS_min)
        test_RSS = test_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # SLF image (shape: [num_sample, K0, K1])
        test_slf_img = testing_data['slf_img'][idx]
        test_slf_img = test_slf_img.reshape(-1, 1, self.args.K0, self.args.K1)

        # SLF image ENR (shape: [num_sample, K0, K1])
        test_slf_img_ENR = testing_data['slf_img_ENR'][idx]
        test_slf_img_ENR = test_slf_img_ENR.reshape(-1, 1, self.args.K0, self.args.K1)

        test_RSS = torch.from_numpy(test_RSS).float()
        test_slf_img = torch.from_numpy(test_slf_img).float()
        test_slf_img_ENR = torch.from_numpy(test_slf_img_ENR).float()

        dataset_test = TensorDataset(test_RSS, test_slf_img_ENR, test_slf_img)

        test_loader = DataLoader(
            dataset_test,
            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 = SLF_Estimator(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)
        # test model
        slf_rmse_model = 0
        slf_rmse_enr = 0
        slf_mae_model = 0
        slf_mae_enr = 0
        labels = list()
        preds_model = list()
        preds_enr = list()
        for batch_idx, (input_rss, slf_enr, target_slf) in enumerate(test_loader):
            input_rss = input_rss.to(self.device)
            slf_enr = slf_enr.to(self.device)
            target_slf = target_slf.to(self.device)
            slf_recon = model(input_rss)
            # test SLF image RMSE
            slf_rmse_model += F.mse_loss(slf_recon, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            slf_rmse_enr += F.mse_loss(slf_enr, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae_model += F.l1_loss(slf_recon, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            slf_mae_enr += F.l1_loss(slf_enr, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            labels.append(target_slf.detach().cpu().numpy())
            preds_model.append(slf_recon.detach().cpu().numpy())
            preds_enr.append(slf_enr.detach().cpu().numpy())

        slf_rmse_model /= len(test_loader.dataset)
        slf_rmse_model = np.sqrt(slf_rmse_model)
        slf_mae_model /= len(test_loader.dataset)
        slf_rmse_enr /= len(test_loader.dataset)
        slf_rmse_enr = np.sqrt(slf_rmse_enr)
        slf_mae_enr /= len(test_loader.dataset)
        print('====> Test Model:')
        print('====> Noise Level: ' + noise_level)
        print('====> Recon SLF rmse: {:.6f}'.format(slf_rmse_model))
        print('====> ENR SLF rmse: {:.6f}'.format(slf_rmse_enr))
        print('====> Recon SLF mae: {:.6f}'.format(slf_mae_model))
        print('====> ENR SLF mae: {:.6f}'.format(slf_mae_enr))

        ## revise part for manual ploting roc curve
        labels = np.concatenate(labels, axis=0)
        np.save('results/compare/ROC_data/labels_' + noise_level + '.npy', labels)
        preds_model = np.concatenate(preds_model, axis=0)
        np.save('results/compare/ROC_data/slf_estimator_preds_'+noise_level+'.npy', preds_model)
        preds_enr = np.concatenate(preds_enr, axis=0)
        np.save('results/compare/ROC_data/enr_preds_' + noise_level + '.npy', preds_enr)

        # # code for auto ploting roc curve
        # labels = np.concatenate(labels, axis=0).reshape(-1)
        # preds1 = np.concatenate(preds1, axis=0).reshape(-1)
        # preds2 = np.concatenate(preds2, axis=0).reshape(-1)
        # labels[labels > 0] = 1
        # save_path1 = self.args.run + 'ENR/' + 'roc_' + noise_level + '.npz'
        # save_path2 = self.path + 'roc_' + noise_level + '.npz'
        # fpr, tpr, thresholds = roc_curve(labels, preds1, pos_label=1)
        # np.savez(save_path1, fpr=fpr, tpr=tpr, thresholds=thresholds)
        # fpr, tpr, thresholds = roc_curve(labels, preds2, pos_label=1)
        # np.savez(save_path2, fpr=fpr, tpr=tpr, thresholds=thresholds)

    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 = SLF_Estimator(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_clean = 'results/compare/YOLO/Origin/images/clean/'+noise_level+'/'
        path_enr = 'results/compare/YOLO/Origin/images/enr/' + noise_level + '/'
        path_model = 'results/compare/YOLO/Origin/images/estimator/' + noise_level + '/'
        if not os.path.exists(path_clean):
            os.makedirs(path_clean)
        if not os.path.exists(path_enr):
            os.makedirs(path_enr)
        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, slf_enr, target_slf) in enumerate(test_loader):
            input_rss = input_rss.to(self.device)
            slf_enr = slf_enr.to(self.device)
            target_slf = target_slf.to(self.device)
            slf_recon = model(input_rss)
            target_slf = target_slf.detach().cpu().numpy()
            slf_enr = slf_enr.detach().cpu().numpy()
            slf_recon = slf_recon.detach().cpu().numpy()
            for idx in range(len(target_slf)):
                clean_img = target_slf[idx, 0, :, :]
                enr_img = slf_enr[idx, 0, :, :]
                # enr_img[enr_img > thr] = 1
                # enr_img[enr_img <= thr] = 0
                estimator_img = slf_recon[idx, 0, :, :]
                # estimator_img[estimator_img > thr] = 1
                # estimator_img[estimator_img <= thr] = 0
                clean_img = cv2.resize(clean_img, (320, 320), interpolation=cv2.INTER_NEAREST)*255
                enr_img = cv2.resize(enr_img, (320, 320), interpolation=cv2.INTER_NEAREST)*255
                estimator_img = cv2.resize(estimator_img, (320, 320), interpolation=cv2.INTER_NEAREST)*255
                cv2.imwrite(path_clean + 'im'+str(idx+batch_idx*self.args.batch_size+1) + '.jpg', clean_img)
                cv2.imwrite(path_enr + 'im' + str(idx+batch_idx*self.args.batch_size+1) +'.jpg', enr_img)
                cv2.imwrite(path_model + 'im' + str(idx+batch_idx*self.args.batch_size+1) +'.jpg', estimator_img)
            # image = cv2.imread(path_enr + 'im' + str(1) +'.jpg', cv2.IMREAD_GRAYSCALE)
            # print(np.amax(image))
            # print(np.shape(image))
            # break


# runner for SLF-VAE
class SLF_VAE_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 + 'SLF_VAE/'
        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))
        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'):
        # Load training data
        training_data = sio.loadmat('./data/training_data.mat')
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        train_RSS = training_data['RSS']
        RSS_min = np.amin(train_RSS, axis=0, keepdims=True)
        RSS_max = np.amax(train_RSS, axis=0, keepdims=True)
        train_RSS = (train_RSS - RSS_min) / (RSS_max - RSS_min)
        train_RSS = train_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # SLF image (shape: [num_sample, K0, K1])
        train_slf_img = training_data['slf_img']
        train_slf_img = train_slf_img.reshape(-1, 1, self.args.K0, self.args.K1)

        # SLF image ENR (shape: [num_sample, K0, K1])
        train_slf_img_ENR = training_data['slf_img_ENR']
        train_slf_img_ENR = train_slf_img_ENR.reshape(-1, 1, self.args.K0, self.args.K1)

        train_RSS = torch.from_numpy(train_RSS).float()
        train_slf_img = torch.from_numpy(train_slf_img).float()
        train_slf_img_ENR = torch.from_numpy(train_slf_img_ENR).float()

        dataset = TensorDataset(train_RSS, train_slf_img_ENR, train_slf_img)
        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))

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

        # SLF image ENR (shape: [num_sample, K0, K1])
        test_slf_img_ENR = testing_data['slf_img_ENR'][idx]
        test_slf_img_ENR = test_slf_img_ENR.reshape(-1, 1, self.args.K0, self.args.K1)

        test_RSS = torch.from_numpy(test_RSS).float()
        test_slf_img = torch.from_numpy(test_slf_img).float()
        test_slf_img_ENR = torch.from_numpy(test_slf_img_ENR).float()

        dataset_test = TensorDataset(test_RSS, test_slf_img_ENR, test_slf_img)

        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)
        print("Data Loaded!")
        return train_loader, val_loader, test_loader

    def train(self, model, train_loader, optimizer, epoch):
        model.train()
        train_loss = 0                  # total loss
        train_recon_loss = 0            # recon loss
        train_kld = 0                   # kld loss
        train_rmse1 = 0                 # rmse for enr
        train_rmse2 = 0                 # rmse for our model
        for batch_idx,  (input_rss, slf_coarse, target_slf) in enumerate(train_loader):
            input_rss = input_rss.to(self.device)
            slf_coarse = slf_coarse.to(self.device)
            target_slf = target_slf.to(self.device)
            optimizer.zero_grad()
            mu, logvar, slf_recon = model.train_model(input_rss, target_slf, slf_coarse)
            recon_loss, kld = model.train_loss(mu, logvar, slf_recon, target_slf)
            loss = self.args.lambda0 * kld + self.args.lambda1 * recon_loss
            train_loss += loss.item()
            train_recon_loss += recon_loss.item()
            train_kld += kld.item()

            mse1 = F.mse_loss(slf_coarse, target_slf, reduction='sum')/(self.args.K0*self.args.K1)
            mse2 = F.mse_loss(slf_recon, target_slf, reduction='sum') / (self.args.K0 * self.args.K1)
            train_rmse1 += mse1.item()
            train_rmse2 += mse2.item()

            loss /= input_rss.size(0)
            loss.backward()
            optimizer.step()
            if batch_idx % self.args.log_interval == 0:
                print(
                    'Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}\tRecon_Loss: {:.6f}\tKLD: {:.6f}'.format(
                        epoch, batch_idx * len(input_rss), len(train_loader.dataset),
                               100. * batch_idx / len(train_loader),
                               loss.item(), recon_loss.item() / len(input_rss),
                               kld.item() / len(input_rss)))

        train_loss /= len(train_loader.dataset)
        train_recon_loss /= len(train_loader.dataset)
        train_kld /= len(train_loader.dataset)
        train_rmse1 /= len(train_loader.dataset)
        train_rmse1 = np.sqrt(train_rmse1)
        train_rmse2 /= len(train_loader.dataset)
        train_rmse2 = np.sqrt(train_rmse2)
        print('====> Epoch: {} Average loss: {:.6f}'.format(epoch, train_loss))
        print('====> Recon loss: {:.6f}'.format(train_recon_loss))
        print('====> KLD: {:.6f}'.format(train_kld))
        print('====> RMSE for ENR: {:.6f}'.format(train_rmse1))
        print('====> RMSE for our model: {:.6f}'.format(train_rmse2))
        return train_loss, train_kld, train_recon_loss, train_rmse1, train_rmse2

    def validate(self, model, val_loader):
        model.eval()
        val_recon_loss = 0           # recon loss
        val_rmse1 = 0                # rmse for ENR
        val_rmse2 = 0                # rmse for our model
        for batch_idx,  (input_rss, slf_coarse, target_slf) in enumerate(val_loader):
            input_rss = input_rss.to(self.device)
            slf_coarse = slf_coarse.to(self.device)
            target_slf = target_slf.to(self.device)
            slf_recon = model.test_model(input_rss, slf_coarse)
            recon_loss = model.test_loss(slf_recon, target_slf)
            val_recon_loss += recon_loss.item()
            mse1 = F.mse_loss(slf_coarse, target_slf, reduction='sum')/(self.args.K0*self.args.K1)
            mse2 = F.mse_loss(slf_recon, target_slf, reduction='sum')/(self.args.K0*self.args.K1)
            val_rmse1 += mse1.item()
            val_rmse2 += mse2.item()
        val_recon_loss /= len(val_loader.dataset)
        val_rmse1 /= len(val_loader.dataset)
        val_rmse1 = np.sqrt(val_rmse1)
        val_rmse2 /= len(val_loader.dataset)
        val_rmse2 = np.sqrt(val_rmse2)
        print('====> Recon loss: {:.6f}'.format(val_recon_loss))
        print('====> RMSE for ENR: {:.6f}'.format(val_rmse1))
        print('====> RMSE for our model: {:.6f}'.format(val_rmse2))
        return val_recon_loss, val_rmse1, val_rmse2

    def test(self, model, test_loader):
        model.eval()
        test_recon_loss = 0  # recon loss
        test_rmse1 = 0  # rmse for ENR
        test_rmse2 = 0  # rmse for our model
        for batch_idx, (input_rss, slf_coarse, target_slf) in enumerate(test_loader):
            input_rss = input_rss.to(self.device)
            slf_coarse = slf_coarse.to(self.device)
            target_slf = target_slf.to(self.device)
            slf_recon = model.test_model(input_rss, slf_coarse)
            recon_loss = model.test_loss(slf_recon, target_slf)
            test_recon_loss += recon_loss.item()
            mse1 = F.mse_loss(slf_coarse, target_slf, reduction='sum') / (self.args.K0 * self.args.K1)
            mse2 = F.mse_loss(slf_recon, target_slf, reduction='sum') / (self.args.K0 * self.args.K1)
            test_rmse1 += mse1.item()
            test_rmse2 += mse2.item()
        test_recon_loss /= len(test_loader.dataset)
        test_rmse1 /= len(test_loader.dataset)
        test_rmse1 = np.sqrt(test_rmse1)
        test_rmse2 /= len(test_loader.dataset)
        test_rmse2 = np.sqrt(test_rmse2)
        print('====> Recon loss: {:.6f}'.format(test_recon_loss))
        print('====> RMSE for ENR: {:.6f}'.format(test_rmse1))
        print('====> RMSE for our model: {:.6f}'.format(test_rmse2))
        return test_recon_loss, test_rmse1, test_rmse2

    def train_save(self):
        model = SLF_VAE(M=self.args.M, P=self.args.P, K=(self.args.K0, self.args.K1))
        model = model.to(self.device)
        train_loader, val_loader, test_loader = self.load_data()
        optimizer = self.get_optimizer(model.parameters())
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.2, patience=8, min_lr=1e-5)
        train_set = []
        val_set = []
        test_set = []
        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) + '.npy'
        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_recon_loss, train_rmse1, train_rmse2 = self.train(model, train_loader,
                                                                                                   optimizer, epoch)
            train_set.append([train_loss, train_kld, train_recon_loss, train_rmse1, train_rmse2])
            print('====> Validation Loss for Epoch {:d}'.format(epoch))
            val_recon_loss, val_rmse1, val_rmse2 = self.validate(model, val_loader)
            val_set.append([val_recon_loss, val_rmse1, val_rmse2])
            scheduler.step(val_recon_loss)
            print('====> Test Loss for Epoch {:d}'.format(epoch))
            test_recon_loss, test_rmse1, test_rmse2 = self.test(model, test_loader)
            test_set.append([test_recon_loss, test_rmse1, test_rmse2])
            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_rmse2 < current_val_slf_loss:
                torch.save(model.state_dict(), model_path)
                current_val_slf_loss = val_rmse2

        train_loss = np.asarray(train_set).reshape(-1, 5)
        val_loss = np.asarray(val_set).reshape(-1, 3)
        test_loss = np.asarray(test_set).reshape(-1, 3)
        np.save(train_loss_path, train_loss)
        np.save(val_loss_path, val_loss)
        np.save(test_loss_path, test_loss)

    def test_model(self, noise_level='all'):
        model = SLF_VAE(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(noise_level)
        # test model
        slf_rmse1 = 0        # rmse for stage 1
        slf_mae1 = 0         # mae for stage 1
        slf_rmse2 = 0        # rmse for stage 2
        slf_mae2 = 0         # mae for stage 2
        for batch_idx, (input_rss, slf_coarse, target_slf) in enumerate(test_loader):
            input_rss = input_rss.to(self.device)
            slf_coarse = slf_coarse.to(self.device)
            target_slf = target_slf.to(self.device)
            slf_recon = model.test_model(input_rss, slf_coarse)
            # plot SLF reconstruction image for image1
            n = min(input_rss.size(0), 8)
            comparison = torch.cat([target_slf.view(target_slf.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    slf_coarse.view(slf_coarse.size(0), 1, self.args.K0, self.args.K1)[:n],
                                    slf_recon.view(slf_recon.size(0), 1, self.args.K0, self.args.K1)[:n]])
            save_image(comparison.cpu(),
                       self.path + 'reconSLF_' + noise_level + '_' + str(batch_idx) + '.png',
                       nrow=n)
            # test SLF image RMSE
            slf_rmse1 += F.mse_loss(slf_coarse, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            slf_rmse2 += F.mse_loss(slf_recon, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae1 += F.l1_loss(slf_coarse, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            slf_mae2 += F.l1_loss(slf_recon, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)

        slf_rmse1 /= len(test_loader.dataset)
        slf_rmse1 = np.sqrt(slf_rmse1)
        slf_rmse2 /= len(test_loader.dataset)
        slf_rmse2 = np.sqrt(slf_rmse2)
        slf_mae1 /= len(test_loader.dataset)
        slf_mae2 /= len(test_loader.dataset)
        print('====> Test Model:')
        print('====> Noise Level: ' + noise_level)
        print('====> ENR SLF rmse: {:.6f}'.format(slf_rmse1))
        print('====> Recon SLF rmse: {:.6f}'.format(slf_rmse2))
        print('====> ENR SLF mae: {:.6f}'.format(slf_mae1))
        print('====> Recon SLF mae: {:.6f}'.format(slf_mae2))

    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_RSS = training_data['RSS']
        RSS_min = np.amin(train_RSS, axis=0, keepdims=True)
        RSS_max = np.amax(train_RSS, axis=0, keepdims=True)

        # Load testing data for statistical results simulation (ROC and mAP)
        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)
        noise_class = testing_data['sig_epsilon_class']
        noise_class = np.squeeze(noise_class)
        idx = (noise_class == 0) | (noise_class == 1) | (noise_class == 2)
        if noise_level == 'all':
            pass
        elif noise_level == 'low':
            idx = (noise_class == 0)
        elif noise_level == 'mid':
            idx = (noise_class == 1)
        elif noise_level == 'high':
            idx = (noise_class == 2)
        else:
            raise NotImplementedError('Noise level {} not understood.'.format(noise_level))
        # normalized noisy RSS measurement input (shape: [num_sample, NxPxP])
        test_RSS = testing_data['RSS'][idx]
        test_RSS = (test_RSS - RSS_min) / (RSS_max - RSS_min)
        test_RSS = test_RSS.reshape(-1, self.args.N, self.args.P, self.args.P)
        # SLF image (shape: [num_sample, K0, K1])
        test_slf_img = testing_data['slf_img'][idx]
        test_slf_img = test_slf_img.reshape(-1, 1, self.args.K0, self.args.K1)

        # SLF image ENR (shape: [num_sample, K0, K1])
        test_slf_img_ENR = testing_data['slf_img_ENR'][idx]
        test_slf_img_ENR = test_slf_img_ENR.reshape(-1, 1, self.args.K0, self.args.K1)

        test_RSS = torch.from_numpy(test_RSS).float()
        test_slf_img = torch.from_numpy(test_slf_img).float()
        test_slf_img_ENR = torch.from_numpy(test_slf_img_ENR).float()

        dataset_test = TensorDataset(test_RSS, test_slf_img_ENR, test_slf_img)

        test_loader = DataLoader(
            dataset_test,
            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 = SLF_VAE(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)
        # test model
        slf_rmse1 = 0        # rmse for stage 1
        slf_mae1 = 0         # mae for stage 1
        slf_rmse2 = 0        # rmse for stage 2
        slf_mae2 = 0         # mae for stage 2
        labels = list()
        preds1 = list()      # for ENR
        preds2 = list()      # for SLF_VAE
        for batch_idx, (input_rss, slf_coarse, target_slf) in enumerate(test_loader):
            input_rss = input_rss.to(self.device)
            slf_coarse = slf_coarse.to(self.device)
            target_slf = target_slf.to(self.device)
            slf_recon = model.test_model(input_rss, slf_coarse)
            # test SLF image RMSE
            slf_rmse1 += F.mse_loss(slf_coarse, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            slf_rmse2 += F.mse_loss(slf_recon, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            # test SLF image MAE
            slf_mae1 += F.l1_loss(slf_coarse, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            slf_mae2 += F.l1_loss(slf_recon, target_slf, reduction='sum').item() / (self.args.K0 * self.args.K1)
            labels.append(target_slf.detach().cpu().numpy())
            preds1.append(slf_coarse.detach().cpu().numpy())
            preds2.append(slf_recon.detach().cpu().numpy())

        slf_rmse1 /= len(test_loader.dataset)
        slf_rmse1 = np.sqrt(slf_rmse1)
        slf_rmse2 /= len(test_loader.dataset)
        slf_rmse2 = np.sqrt(slf_rmse2)
        slf_mae1 /= len(test_loader.dataset)
        slf_mae2 /= len(test_loader.dataset)
        print('====> Test Model:')
        print('====> Noise Level: ' + noise_level)
        print('====> ENR SLF rmse: {:.6f}'.format(slf_rmse1))
        print('====> Recon SLF rmse: {:.6f}'.format(slf_rmse2))
        print('====> ENR SLF mae: {:.6f}'.format(slf_mae1))
        print('====> Recon SLF mae: {:.6f}'.format(slf_mae2))

        ## revise part for manual ploting roc curve
        preds_model = np.concatenate(preds2, axis=0)
        np.save('results/compare/ROC_data/slf_vae_preds_' + noise_level + '.npy', preds_model)