import os
import sys
import logging
import time
import argparse
import numpy as np
from collections import OrderedDict

import scripts.options as option
import utils.util as util
from data.util import bgr2ycbcr
from data import create_dataset, create_dataloader
from models import create_model

# options
parser = argparse.ArgumentParser()
parser.add_argument('-opt', type=str, required=True, help='Path to options JSON file.')
opt = option.parse(parser.parse_args().opt, is_train=False)
util.mkdirs((path for key, path in opt['path'].items() if not key == 'pretrain_model_G'))
opt = option.dict_to_nonedict(opt)

util.setup_logger(None, opt['path']['log'], 'test.log', level=logging.INFO, screen=True)
logger = logging.getLogger('base')
logger.info(option.dict2str(opt))
# Create test dataset and dataloader
test_loaders = []
for phase, dataset_opt in sorted(opt['datasets'].items()):
    test_set = create_dataset(dataset_opt)
    test_loader = create_dataloader(test_set, dataset_opt)
    logger.info('Number of test images in [{:s}]: {:d}'.format(dataset_opt['name'], len(test_set)))
    test_loaders.append(test_loader)

# Create model
model = create_model(opt)

for test_loader in test_loaders:
    test_set_name = test_loader.dataset.opt['name']
    logger.info('\nTesting [{:s}]...'.format(test_set_name))
    test_start_time = time.time()
    dataset_dir = os.path.join(opt['path']['results_root'], test_set_name)
    util.mkdir(dataset_dir)

    test_results = OrderedDict()
    test_results['psnr'] = []
    test_results['ssim'] = []
    test_results['psnr_y'] = []
    test_results['ssim_y'] = []

    for data in test_loader:
        need_GT = False if test_loader.dataset.opt['dataroot_GT'] is None else True
        # need_GT = True

        model.feed_data_specular(data, need_GT=need_GT)
        if opt["image_type"] == "exr":
            y = data["x_offset"]
            x = data["y_offset"]
        img_path = data['NOISY_path'][0]
        img_name = os.path.splitext(os.path.basename(img_path))[0]
        start = time.time()
        model.test()  # test
        end = time.time()
        print("Time elapsed... %f "%(end - start))
        visuals = model.get_current_visuals(need_GT=need_GT)

        denoised_img = util.tensor2img(visuals['DENOISED'])  # uint8
        noisy_img = util.tensor2img(visuals['NOISY'])
        gt_img = util.tensor2img(visuals['GT'])  # uint8

        # save images
        suffix = opt['suffix']
        if suffix ==None:
            suffix = ""
            
        save_DENOISED_img_path = os.path.join(dataset_dir, img_name + suffix + '_1denoised.png')
        save_NOISY_img_path = os.path.join(dataset_dir, img_name + suffix + '_0noisy.png')
        save_GT_img_path = os.path.join(dataset_dir, img_name + suffix + '_2gt.png')
  
        
        # calculate PSNR and SSIM
        if need_GT:
            # gt_img = util.tensor2img(visuals['GT'])
            gt_img = gt_img / 255.
            denoised_img = denoised_img / 255.

            crop_border = test_loader.dataset.opt['scale']
            cropped_denoised_img = denoised_img#[crop_border:-crop_border, crop_border:-crop_border, :]
            cropped_gt_img = gt_img#[crop_border:-crop_border, crop_border:-crop_border, :]

            psnr = util.calculate_psnr(cropped_denoised_img * 255, cropped_gt_img * 255)
            ssim = util.calculate_ssim(cropped_denoised_img * 255, cropped_gt_img * 255)
            test_results['psnr'].append(psnr)
            test_results['ssim'].append(ssim)

            if gt_img.shape[2] == 3:  # RGB image
                denoised_img_y = bgr2ycbcr(denoised_img, only_y=True)
                gt_img_y = bgr2ycbcr(gt_img, only_y=True)
                cropped_denoised_img_y = denoised_img_y[crop_border:-crop_border, crop_border:-crop_border]
                cropped_gt_img_y = gt_img_y[crop_border:-crop_border, crop_border:-crop_border]
                psnr_y = util.calculate_psnr(cropped_denoised_img_y * 255, cropped_gt_img_y * 255)
                ssim_y = util.calculate_ssim(cropped_denoised_img_y * 255, cropped_gt_img_y * 255)
                test_results['psnr_y'].append(psnr_y)
                test_results['ssim_y'].append(ssim_y)
                logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}; PSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}.'\
                    .format(img_name, psnr, ssim, psnr_y, ssim_y))
            else:
                logger.info('{:20s} - PSNR: {:.6f} dB; SSIM: {:.6f}.'.format(img_name, psnr, ssim))
        else:
            logger.info(img_name)


        if opt["image_type"] == "exr":
            denoised_exr = util.tensor2exr(visuals['DENOISED'])  # uint8
            noisy_exr = util.tensor2exr(visuals['NOISY'])
            gt_exr = util.tensor2exr(visuals['GT'])  # uint8

            save_DENOISED_img_path = os.path.join(dataset_dir, img_name + suffix + '_1denoised.exr')
            save_NOISY_img_path = os.path.join(dataset_dir, img_name + suffix + '_0noisy.exr')
            save_GT_img_path = os.path.join(dataset_dir, img_name + suffix + '_2gt.exr') 
  
            util.saveEXRfromMatrix(save_DENOISED_img_path, denoised_exr, (x, y)) 
            util.saveEXRfromMatrix(save_NOISY_img_path, noisy_exr, (x, y))  
            util.saveEXRfromMatrix(save_GT_img_path, gt_exr, (x, y))

    if need_GT:  # metrics
        # Average PSNR/SSIM results
        ave_psnr = sum(test_results['psnr']) / len(test_results['psnr'])
        ave_ssim = sum(test_results['ssim']) / len(test_results['ssim'])
        logger.info('----Average PSNR/SSIM results for {}----\n\tPSNR: {:.6f} dB; SSIM: {:.6f}\n'\
                .format(test_set_name, ave_psnr, ave_ssim))
        # if test_results['psnr_y'] and test_results['ssim_y']:
        #     ave_psnr_y = sum(test_results['psnr_y']) / len(test_results['psnr_y'])
        #     ave_ssim_y = sum(test_results['ssim_y']) / len(test_results['ssim_y'])
        #     logger.info('----Y channel, average PSNR/SSIM----\n\tPSNR_Y: {:.6f} dB; SSIM_Y: {:.6f}\n'\
        #         .format(ave_psnr_y, ave_ssim_y))

[((16, 9, 16, 34), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ({}, {}), '()', False, 'import argparse\n'), ((20, 6, 20, 34), 'scripts.options.dict_to_nonedict', 'option.dict_to_nonedict', ({(20, 30, 20, 33): 'opt'}, {}), '(opt)', True, 'import scripts.options as option\n'), ((22, 0, 22, 88), 'utils.util.setup_logger', 'util.setup_logger', (), '', True, 'import utils.util as util\n'), ((23, 9, 23, 34), 'logging.getLogger', 'logging.getLogger', ({(23, 27, 23, 33): '"""base"""'}, {}), "('base')", False, 'import logging\n'), ((34, 8, 34, 25), 'models.create_model', 'create_model', ({(34, 21, 34, 24): 'opt'}, {}), '(opt)', False, 'from models import create_model\n'), ((24, 12, 24, 32), 'scripts.options.dict2str', 'option.dict2str', ({(24, 28, 24, 31): 'opt'}, {}), '(opt)', True, 'import scripts.options as option\n'), ((28, 15, 28, 42), 'data.create_dataset', 'create_dataset', ({(28, 30, 28, 41): 'dataset_opt'}, {}), '(dataset_opt)', False, 'from data import create_dataset, create_dataloader\n'), ((29, 18, 29, 58), 'data.create_dataloader', 'create_dataloader', ({(29, 36, 29, 44): 'test_set', (29, 46, 29, 57): 'dataset_opt'}, {}), '(test_set, dataset_opt)', False, 'from data import create_dataset, create_dataloader\n'), ((39, 22, 39, 33), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((40, 18, 40, 74), 'os.path.join', 'os.path.join', ({(40, 31, 40, 58): "opt['path']['results_root']", (40, 60, 40, 73): 'test_set_name'}, {}), "(opt['path']['results_root'], test_set_name)", False, 'import os\n'), ((41, 4, 41, 27), 'utils.util.mkdir', 'util.mkdir', ({(41, 15, 41, 26): 'dataset_dir'}, {}), '(dataset_dir)', True, 'import utils.util as util\n'), ((43, 19, 43, 32), 'collections.OrderedDict', 'OrderedDict', ({}, {}), '()', False, 'from collections import OrderedDict\n'), ((59, 16, 59, 27), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((61, 14, 61, 25), 'time.time', 'time.time', ({}, {}), '()', False, 'import time\n'), ((65, 23, 65, 59), 'utils.util.tensor2img', 'util.tensor2img', ({(65, 39, 65, 58): "visuals['DENOISED']"}, {}), "(visuals['DENOISED'])", True, 'import utils.util as util\n'), ((66, 20, 66, 53), 'utils.util.tensor2img', 'util.tensor2img', ({(66, 36, 66, 52): "visuals['NOISY']"}, {}), "(visuals['NOISY'])", True, 'import utils.util as util\n'), ((67, 17, 67, 47), 'utils.util.tensor2img', 'util.tensor2img', ({(67, 33, 67, 46): "visuals['GT']"}, {}), "(visuals['GT'])", True, 'import utils.util as util\n'), ((74, 33, 74, 96), 'os.path.join', 'os.path.join', ({(74, 46, 74, 57): 'dataset_dir', (74, 59, 74, 95): "img_name + suffix + '_1denoised.png'"}, {}), "(dataset_dir, img_name + suffix + '_1denoised.png')", False, 'import os\n'), ((75, 30, 75, 90), 'os.path.join', 'os.path.join', ({(75, 43, 75, 54): 'dataset_dir', (75, 56, 75, 89): "img_name + suffix + '_0noisy.png'"}, {}), "(dataset_dir, img_name + suffix + '_0noisy.png')", False, 'import os\n'), ((76, 27, 76, 84), 'os.path.join', 'os.path.join', ({(76, 40, 76, 51): 'dataset_dir', (76, 53, 76, 83): "img_name + suffix + '_2gt.png'"}, {}), "(dataset_dir, img_name + suffix + '_2gt.png')", False, 'import os\n'), ((89, 19, 89, 88), 'utils.util.calculate_psnr', 'util.calculate_psnr', ({(89, 39, 89, 65): 'cropped_denoised_img * 255', (89, 67, 89, 87): 'cropped_gt_img * 255'}, {}), '(cropped_denoised_img * 255, cropped_gt_img * 255)', True, 'import utils.util as util\n'), ((90, 19, 90, 88), 'utils.util.calculate_ssim', 'util.calculate_ssim', ({(90, 39, 90, 65): 'cropped_denoised_img * 255', (90, 67, 90, 87): 'cropped_gt_img * 255'}, {}), '(cropped_denoised_img * 255, cropped_gt_img * 255)', True, 'import utils.util as util\n'), ((112, 27, 112, 63), 'utils.util.tensor2exr', 'util.tensor2exr', ({(112, 43, 112, 62): "visuals['DENOISED']"}, {}), "(visuals['DENOISED'])", True, 'import utils.util as util\n'), ((113, 24, 113, 57), 'utils.util.tensor2exr', 'util.tensor2exr', ({(113, 40, 113, 56): "visuals['NOISY']"}, {}), "(visuals['NOISY'])", True, 'import utils.util as util\n'), ((114, 21, 114, 51), 'utils.util.tensor2exr', 'util.tensor2exr', ({(114, 37, 114, 50): "visuals['GT']"}, {}), "(visuals['GT'])", True, 'import utils.util as util\n'), ((116, 37, 116, 100), 'os.path.join', 'os.path.join', ({(116, 50, 116, 61): 'dataset_dir', (116, 63, 116, 99): "img_name + suffix + '_1denoised.exr'"}, {}), "(dataset_dir, img_name + suffix + '_1denoised.exr')", False, 'import os\n'), ((117, 34, 117, 94), 'os.path.join', 'os.path.join', ({(117, 47, 117, 58): 'dataset_dir', (117, 60, 117, 93): "img_name + suffix + '_0noisy.exr'"}, {}), "(dataset_dir, img_name + suffix + '_0noisy.exr')", False, 'import os\n'), ((118, 31, 118, 88), 'os.path.join', 'os.path.join', ({(118, 44, 118, 55): 'dataset_dir', (118, 57, 118, 87): "img_name + suffix + '_2gt.exr'"}, {}), "(dataset_dir, img_name + suffix + '_2gt.exr')", False, 'import os\n'), ((120, 12, 120, 80), 'utils.util.saveEXRfromMatrix', 'util.saveEXRfromMatrix', ({(120, 35, 120, 57): 'save_DENOISED_img_path', (120, 59, 120, 71): 'denoised_exr', (120, 73, 120, 79): '(x, y)'}, {}), '(save_DENOISED_img_path, denoised_exr, (x, y))', True, 'import utils.util as util\n'), ((121, 12, 121, 74), 'utils.util.saveEXRfromMatrix', 'util.saveEXRfromMatrix', ({(121, 35, 121, 54): 'save_NOISY_img_path', (121, 56, 121, 65): 'noisy_exr', (121, 67, 121, 73): '(x, y)'}, {}), '(save_NOISY_img_path, noisy_exr, (x, y))', True, 'import utils.util as util\n'), ((122, 12, 122, 68), 'utils.util.saveEXRfromMatrix', 'util.saveEXRfromMatrix', ({(122, 35, 122, 51): 'save_GT_img_path', (122, 53, 122, 59): 'gt_exr', (122, 61, 122, 67): '(x, y)'}, {}), '(save_GT_img_path, gt_exr, (x, y))', True, 'import utils.util as util\n'), ((58, 36, 58, 62), 'os.path.basename', 'os.path.basename', ({(58, 53, 58, 61): 'img_path'}, {}), '(img_path)', False, 'import os\n'), ((95, 33, 95, 69), 'data.util.bgr2ycbcr', 'bgr2ycbcr', (), '', False, 'from data.util import bgr2ycbcr\n'), ((96, 27, 96, 57), 'data.util.bgr2ycbcr', 'bgr2ycbcr', (), '', False, 'from data.util import bgr2ycbcr\n'), ((99, 25, 99, 98), 'utils.util.calculate_psnr', 'util.calculate_psnr', ({(99, 45, 99, 73): 'cropped_denoised_img_y * 255', (99, 75, 99, 97): 'cropped_gt_img_y * 255'}, {}), '(cropped_denoised_img_y * 255, cropped_gt_img_y * 255)', True, 'import utils.util as util\n'), ((100, 25, 100, 98), 'utils.util.calculate_ssim', 'util.calculate_ssim', ({(100, 45, 100, 73): 'cropped_denoised_img_y * 255', (100, 75, 100, 97): 'cropped_gt_img_y * 255'}, {}), '(cropped_denoised_img_y * 255, cropped_gt_img_y * 255)', True, 'import utils.util as util\n')]