import os
import torch
import argparse
import numpy as np
from utils1.log import get_logger
from utils1.datasets import PIPAL_demo, five_point_crop
from models.build import build_model
from torch.utils.data import DataLoader
from scipy.stats import spearmanr, pearsonr

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')

    # dir setting
    parser.add_argument('--ref_path', type=str, default='./datasets/PIPAL/test/ref/')
    parser.add_argument('--dis_path', type=str, default='./datasets/PIPAL/test/dis')
    parser.add_argument('--txt_file_name', type=str, default='./datasets/PIPAL/test/Test_Label/label.txt')
    parser.add_argument('--log_dir', type=str, default='./log')
    parser.add_argument('--model_dir', type=str, default='pretrained')

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

    # mode
    parser.add_argument('--mode', type=str, default='val')

    args = parser.parse_args()
    return args


def update_state_dict(model, pretrained_dict):
    own_state_dict = model.state_dict()
    # filter parameters
    pretrained_dict = {name: param for name, param in pretrained_dict.items() if name in own_state_dict and 'mask' not in name}
    own_state_dict.update(pretrained_dict)
    model.load_state_dict(own_state_dict)


def test(model, logger, test_loader, device, dataset, iter):
    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']
            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)
        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 = ['pipal']

    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}')

    model_file = os.path.join(args.log_dir, args.model_dir, 'train_kadid10k/_iter_45k.tar')
    checkpoint = torch.load(model_file, map_location=args.device)
    update_state_dict(model, checkpoint['model_state_dict'])

    for dataset in datasets:
        logger.debug(f'Current dataset: {dataset}')
        testset = PIPAL_demo(args.ref_path, args.dis_path, args.txt_file_name)
        test_loader = DataLoader(testset, batch_size=1, num_workers=4, shuffle=False)
        srocc, plcc = test(model, logger, test_loader, args.device, dataset, iter=45)
        logger.debug(f'srocc: {srocc}, plcc: {plcc:.4f}')
    logger.debug('-----------------------------------------------------------\n')
