# config
import sys
import warnings
import time
import numpy as np
import torch
from config import config
from eval.Estimator import Estimator
from net.networks import *
from options.test_options import TestOptions
from Dataset.DatasetConstructor import TrainDatasetConstructor,EvalDatasetConstructor

opt = TestOptions().parse()
opt.nThreads = 1   # test code only supports nThreads = 1
opt.batch_size = 1  # test code only supports batchSize = 1
opt.is_flip = 0  # no flip

#test_model_name = 'output/HRNet_relu_aspp_b16/JSTL/MAE_54.1_MSE_116.99_mae_58.27_9.66_93.87_mse_97.44_22.59_168.64_Ep_374.pth'
#test_model_name = 'output/HRNet_relu_aspp_b16/JSTL/MAE_56.1_MSE_123.66_mae_62.8_11.21_94.91_mse_99.34_18.25_180.0_Ep_220.pth'
#test_model_name = 'output/HRNet_relu_aspp_b16/JSTL/MAE_54.89_MSE_122.7_mae_59.64_10.16_94.61_mse_99.23_17.83_178.42_Ep_345.pth'
#test_model_name = 'output/HRNet_relu_aspp_b16/JSTL/MAE_55.02_MSE_117.63_mae_62.98_10.75_92.55_mse_101.18_16.22_169.24_Ep_271.pth'




test_model_name = 'output/HRNet_relu_aspp_fixEvalBug/JSTL_SH_LgQF_NU_JU/MAE_69.45_MSE_288.17_mae_66.05_12.21_104.91_87.49_68.11_mse_117.92_19.51_185.47_494.58_258.65_Ep_26.pth'
save_model_MAE  = '69.45_p1'


# Mainly get settings for specific datasets
setting = config(opt)

# Data loaders
eval_dataset = EvalDatasetConstructor(
    setting.datasets_com,
    setting.eval_img_path,
    setting.eval_gt_map_path,
    mode=setting.mode,
    dataset_name=setting.dataset_name,
    device=setting.device)
eval_loader = torch.utils.data.DataLoader(dataset=eval_dataset, batch_size=1)

# model construct
net = define_net(opt.net_name)
net = init_net(net, gpu_ids=opt.gpu_ids)
net.module.load_state_dict(torch.load(test_model_name, map_location=str(setting.device)))
criterion = torch.nn.MSELoss(reduction='sum').to(setting.device)
estimator = Estimator(setting, eval_loader, criterion=criterion)


validate_MAE, validate_RMSE, validate_loss, time_cost = estimator.evaluate(net, weight_scale=opt.weight_scale, save_model_MAE=save_model_MAE) # I add weight_scale here
sys.stdout.write('loss = {}, eval_mae = {}, eval_rmse = {}, time cost eval = {}s\n'
                .format(validate_loss, validate_MAE, validate_RMSE, time_cost))
sys.stdout.flush()

# config
import sys
import warnings
import time
import numpy as np
import torch
from config import config
from eval.Estimator import Estimator
from net.networks import *
from options.test_options import TestOptions
from Dataset.DatasetConstructor import TrainDatasetConstructor,EvalDatasetConstructor

opt = TestOptions().parse()
opt.nThreads = 1   # test code only supports nThreads = 1
opt.batch_size = 1  # test code only supports batchSize = 1
opt.is_flip = 0  # no flip

#test_model_name = 'output/HRNet_relu_aspp_b16/JSTL/MAE_54.1_MSE_116.99_mae_58.27_9.66_93.87_mse_97.44_22.59_168.64_Ep_374.pth'
#test_model_name = 'output/HRNet_relu_aspp_b16/JSTL/MAE_56.1_MSE_123.66_mae_62.8_11.21_94.91_mse_99.34_18.25_180.0_Ep_220.pth'
#test_model_name = 'output/HRNet_relu_aspp_b16/JSTL/MAE_54.89_MSE_122.7_mae_59.64_10.16_94.61_mse_99.23_17.83_178.42_Ep_345.pth'
#test_model_name = 'output/HRNet_relu_aspp_b16/JSTL/MAE_55.02_MSE_117.63_mae_62.98_10.75_92.55_mse_101.18_16.22_169.24_Ep_271.pth'




test_model_name = 'output/HRNet_relu_aspp_fixEvalBug/JSTL_SH_LgQF_NU_JU/MAE_69.45_MSE_288.17_mae_66.05_12.21_104.91_87.49_68.11_mse_117.92_19.51_185.47_494.58_258.65_Ep_26.pth'
save_model_MAE  = '69.45_p1'


# Mainly get settings for specific datasets
setting = config(opt)

# Data loaders
eval_dataset = EvalDatasetConstructor(
    setting.datasets_com,
    setting.eval_img_path,
    setting.eval_gt_map_path,
    mode=setting.mode,
    dataset_name=setting.dataset_name,
    device=setting.device)
eval_loader = torch.utils.data.DataLoader(dataset=eval_dataset, batch_size=1)

# model construct
net = define_net(opt.net_name)
net = init_net(net, gpu_ids=opt.gpu_ids)
net.module.load_state_dict(torch.load(test_model_name, map_location=str(setting.device)))
criterion = torch.nn.MSELoss(reduction='sum').to(setting.device)
estimator = Estimator(setting, eval_loader, criterion=criterion)


validate_MAE, validate_RMSE, validate_loss, time_cost = estimator.evaluate(net, weight_scale=opt.weight_scale, save_model_MAE=save_model_MAE) # I add weight_scale here
sys.stdout.write('loss = {}, eval_mae = {}, eval_rmse = {}, time cost eval = {}s\n'
                .format(validate_loss, validate_MAE, validate_RMSE, time_cost))
sys.stdout.flush()

# config # test code only supports nThreads = 1 # test code only supports batchSize = 1 # no flip #test_model_name = 'output/HRNet_relu_aspp_b16/JSTL/MAE_54.1_MSE_116.99_mae_58.27_9.66_93.87_mse_97.44_22.59_168.64_Ep_374.pth' #test_model_name = 'output/HRNet_relu_aspp_b16/JSTL/MAE_56.1_MSE_123.66_mae_62.8_11.21_94.91_mse_99.34_18.25_180.0_Ep_220.pth' #test_model_name = 'output/HRNet_relu_aspp_b16/JSTL/MAE_54.89_MSE_122.7_mae_59.64_10.16_94.61_mse_99.23_17.83_178.42_Ep_345.pth' #test_model_name = 'output/HRNet_relu_aspp_b16/JSTL/MAE_55.02_MSE_117.63_mae_62.98_10.75_92.55_mse_101.18_16.22_169.24_Ep_271.pth' # Mainly get settings for specific datasets # Data loaders # model construct # I add weight_scale here