import os
import torch
import argparse
import numpy as np

from utils1.log import get_logger
from utils1.datasets import five_point_crop, TestSet
from models.build import build_model
from torch.utils.data import DataLoader
from scipy.stats import spearmanr, pearsonr

from scipy.io import savemat

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"


def get_args():
    parser = argparse.ArgumentParser('val processing')
    parser.add_argument('--arch', type=str, default='M2DANet')
    parser.add_argument('--use_checkpoint', action='store_true')
    parser.add_argument('--crop_size', type=int, default=224)
    parser.add_argument('--n_ensemble', type=int, default=5, help='ensemble ways of validation')

    parser.add_argument('--dataset_dir', type=str, default='./datasets/')
    parser.add_argument('--log_dir', type=str, default='./log')
    # 要改
    parser.add_argument('--model_dir', type=str, default='pretrained')
    parser.add_argument('--train_data', type=str, default='kadid10k', help='kadid10k | pipal')
    # mode
    parser.add_argument('--mode', type=str, default='val')

    # device
    parser.add_argument('--device', type=str, default='cuda:0')
    parser.add_argument('--gpu_id', type=str, default='0')

    args = parser.parse_args()
    return args


def test(model, logger, test_loader, device, dataset, iter):
    if dataset == 'LIVE':
        test_img_num = 779
    elif dataset == 'CSIQ':
        test_img_num = 866
    else:
        test_img_num = 3000
    test_img_num = 10  # for sample test only
    name = np.empty((test_img_num, 1), dtype=object)
    model.eval()
    with torch.no_grad():
        pred_epoch = []
        labels_epoch = []
        for idx, sample in enumerate(test_loader):
            d_img_org = sample['d_img_org'].to(device)
            r_img_org = sample['r_img_org'].to(device)
            labels = sample['score'].to(device)
            labels = torch.squeeze(labels.type(torch.FloatTensor)).to(device)
            d_img_name = sample['d_img_name']
            name[idx, 0] = d_img_name[0]
            pred = 0
            for i in range(args.n_ensemble):
                b, c, h, w = r_img_org.size()
                # 大于9：随机剪裁
                if args.n_ensemble > 9:
                    new_h = args.crop_size
                    new_w = args.crop_size
                    top = np.random.randint(0, h - new_h)
                    left = np.random.randint(0, w - new_w)
                    r_img = r_img_org[:, :, top: top + new_h, left: left + new_w]
                    d_img = d_img_org[:, :, top: top + new_h, left: left + new_w]
                # 等于1：不剪裁
                elif args.n_ensemble == 1:
                    r_img = r_img_org
                    d_img = d_img_org
                else:
                    # 小于9：五点或者九点剪裁
                    d_img, r_img = five_point_crop(i, d_img=d_img_org, r_img=r_img_org, crop_size=224)
                d_img = d_img.to(device)
                r_img = r_img.to(device)
                pred += model(d_img, r_img)
            pred /= args.n_ensemble

            pred_batch_numpy = pred.squeeze().cpu().numpy()
            labels_batch_numpy = labels.squeeze().cpu().numpy()
            print('iter:', iter, ' d_img_name:', d_img_name[0],
                  ' pred:', str(pred_batch_numpy)[:6], ' gt:', str(labels_batch_numpy)[:6])

            pred_epoch.append(pred_batch_numpy)
            labels_epoch.append(labels_batch_numpy)

        pred_epoch = np.array(pred_epoch)
        labels_epoch = np.array(labels_epoch)
        save_dir = os.path.join(args.log_dir, args.mode, 'result', 'train_' + args.train_data)
        os.makedirs(save_dir, exist_ok=True)
        mat_file = os.path.join(save_dir, 'res_' + dataset + '_' + str(int(iter)) + 'k.mat')
        savemat(mat_file, {'predict': pred_epoch, 'label': labels_epoch, 'name': name})

        rho_s, _ = spearmanr(pred_epoch, labels_epoch)
        rho_p, _ = pearsonr(pred_epoch, labels_epoch)

        logger.debug(f'srcc -- {rho_s:.5f},  plcc -- {rho_p:.5f}')

        return rho_s, rho_p


if __name__ == '__main__':
    args = get_args()
    logger = get_logger(args.log_dir, args.mode, args)

    datasets = ['LIVE', 'CSIQ', 'TID2013']

    if args.train_data == 'kadid10k':
        start_iter = 45000
        end_iter = 45000
        step = 5000
    elif args.train_data == 'pipal':
        start_iter = 55000
        end_iter = 55000
        step = 5000
    else:
        raise NotImplementedError('The %s dataset was not trained' % args.train_data)

    iters = np.arange(start_iter, end_iter + 1, step)
    model = build_model(args.arch, args.gpu_id).to(args.device)
    torch.backends.cudnn.benchmark = True
    total = sum([param.nelement() for param in model.parameters() if param.requires_grad])
    logger.debug("Number of parameter: %.2fM" % (total / 1e6))
    logger.debug(f'Model type: {args.arch}')

    srcc = {dataset: np.zeros(len(iters)) for dataset in datasets}
    plcc = {dataset: np.zeros(len(iters)) for dataset in datasets}

    for idx, iteration in enumerate(iters):
        model_file = os.path.join(args.log_dir, args.model_dir, 'train_' + args.train_data, f'_iter_{int(iteration/1000)}k.tar')
        logger.info(f'Load model: {model_file}...')
        checkpoint = torch.load(model_file, map_location=args.device)
        model.load_state_dict(checkpoint['model_state_dict'], strict=False)
        for dataset in datasets:
            logger.debug(f'Current dataset: {dataset}')
            args.testset = dataset
            testset = TestSet(args.dataset_dir, args.testset)
            test_loader = DataLoader(testset, batch_size=1, num_workers=4, shuffle=False)
            current_srcc, current_plcc = test(model, logger, test_loader, args.device, dataset, iter=iteration / 1000)
            srcc[dataset][idx] = current_srcc
            plcc[dataset][idx] = current_plcc
        logger.debug('-----------------------------------------------------------\n')

    logger.debug('===========================================================')
    for dataset in datasets:
        logger.info(f'Dataset: {dataset}')
        logger.info(f'Best SRCC: {np.max(srcc[dataset]):.2f}, iter: {iters[np.argmax(srcc[dataset])]}')
        logger.info(f'Best PLCC: {np.max(plcc[dataset]):.4f}, iter: {iters[np.argmax(plcc[dataset])]}')
    logger.debug('===========================================================\n')

