import argparse

import numpy as np
import torch
import torch.utils.data
import torchvision.datasets as datasets
import torchvision.transforms as transforms
from networks.resnet import resnet50
from sklearn.metrics import accuracy_score, average_precision_score
from tqdm import tqdm

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('-d', '--dir', nargs='+', type=str, default='examples/realfakedir')
parser.add_argument(
    '-m', '--model_path', type=str, default='weights/blur_jpg_prob0.5.pth'
)
parser.add_argument('-b', '--batch_size', type=int, default=32)
parser.add_argument('-j', '--workers', type=int, default=4, help='number of workers')
parser.add_argument(
    '-c',
    '--crop',
    type=int,
    default=None,
    help='by default, do not crop. specify crop size',
)
parser.add_argument(
    '--use_cpu', action='store_true', help='uses gpu by default, turn on to use cpu'
)
parser.add_argument(
    '--size_only', action='store_true', help='only look at sizes of images in dataset'
)

opt = parser.parse_args()

# Load model
if not opt.size_only:
    model = resnet50(num_classes=1)
    if opt.model_path is not None:
        state_dict = torch.load(opt.model_path, map_location='cpu')
    model.load_state_dict(state_dict['model'])
    model.eval()
    if not opt.use_cpu:
        model.cuda()

# Transform
trans_init = []
if opt.crop is not None:
    trans_init = [
        transforms.CenterCrop(opt.crop),
    ]
    print('Cropping to [%i]' % opt.crop)
else:
    print('Not cropping')
trans = transforms.Compose(
    trans_init
    + [
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ]
)

# Dataset loader
if type(opt.dir) == str:
    opt.dir = [
        opt.dir,
    ]

print('Loading [%i] datasets' % len(opt.dir))
data_loaders = []
for dir in opt.dir:
    dataset = datasets.ImageFolder(dir, transform=trans)
    data_loaders += [
        torch.utils.data.DataLoader(
            dataset, batch_size=opt.batch_size, shuffle=False, num_workers=opt.workers
        ),
    ]

y_true, y_pred = [], []
Hs, Ws = [], []
with torch.no_grad():
    for data_loader in data_loaders:
        for data, label in tqdm(data_loader):
            # for data, label in data_loader:
            Hs.append(data.shape[2])
            Ws.append(data.shape[3])

            y_true.extend(label.flatten().tolist())
            if not opt.size_only:
                if not opt.use_cpu:
                    data = data.cuda()
                y_pred.extend(model(data).sigmoid().flatten().tolist())

Hs, Ws = np.array(Hs), np.array(Ws)
y_true, y_pred = np.array(y_true), np.array(y_pred)

print(
    'Average sizes: [{:2.2f}+/-{:2.2f}] x [{:2.2f}+/-{:2.2f}] = [{:2.2f}+/-{:2.2f} Mpix]'.format(
        np.mean(Hs),
        np.std(Hs),
        np.mean(Ws),
        np.std(Ws),
        np.mean(Hs * Ws) / 1e6,
        np.std(Hs * Ws) / 1e6,
    )
)
print('Num reals: {}, Num fakes: {}'.format(np.sum(1 - y_true), np.sum(y_true)))

if not opt.size_only:
    r_acc = accuracy_score(y_true[y_true == 0], y_pred[y_true == 0] > 0.5)
    f_acc = accuracy_score(y_true[y_true == 1], y_pred[y_true == 1] > 0.5)
    acc = accuracy_score(y_true, y_pred > 0.5)
    ap = average_precision_score(y_true, y_pred)

    print(
        'AP: {:2.2f}, Acc: {:2.2f}, Acc (real): {:2.2f}, Acc (fake): {:2.2f}'.format(
            ap * 100.0, acc * 100.0, r_acc * 100.0, f_acc * 100.0
        )
    )