from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import _init_paths

import os
import cv2
import numpy as np
from progress.bar import Bar
import torch

from utils.utils import AverageMeter
from datasets.dataset_factory import dataset_factory
from detectors.detector_factory import detector_factory


class PrefetchDataset(torch.utils.data.Dataset):
    def __init__(self, opt, dataset, pre_process_func):
        self.images = dataset.images
        self.load_image_func = dataset.coco.loadImgs
        self.img_dir = dataset.img_dir
        self.pre_process_func = pre_process_func
        self.opt = opt

    def __getitem__(self, index):
        img_id = self.images[index]
        img_info = self.load_image_func(ids=[img_id])[0]
        img_path = os.path.join(self.img_dir, img_info['file_name'])
        image = cv2.imread(img_path)
        images, meta = {}, {}
        for scale in self.opt.test_scales:
            if self.opt.task == 'ddd':
                images[scale], meta[scale] = self.pre_process_func(
                    image, scale, img_info['calib'])
            else:
                images[scale], meta[scale] = self.pre_process_func(image, scale)
        return img_id, {'images': images, 'image': image, 'meta': meta}

    def __len__(self):
        return len(self.images)


class Evaluator(object):
    def __init__(self, opt, model):
        self.opt = opt
        if isinstance(model, torch.nn.DataParallel):
            model = model.module
        self.model = model
        Dataset = dataset_factory[opt.dataset]
        Detector = detector_factory[opt.task]
        self.dataset = Dataset(opt, 'val')
        self.detector = Detector(opt, model)
        self.data_loader = torch.utils.data.DataLoader(
            PrefetchDataset(opt, self.dataset, self.detector.pre_process),
            batch_size=1, shuffle=False, num_workers=1, pin_memory=True)

    def eval(self):
        self.model.eval()
        # torch.cuda.empty_cache()
        results = {}
        num_iters = len(self.dataset)
        bar = Bar('{}/{}'.format(self.opt.task, self.opt.exp_id), max=num_iters)
        # time_stats = ['tot', 'load', 'pre', 'net', 'dec', 'post', 'merge']
        # avg_time_stats = {t: AverageMeter() for t in time_stats}
        for ind, (img_id, pre_processed_images) in enumerate(self.data_loader):
            ret = self.detector.run(pre_processed_images)
            results[img_id.numpy().astype(np.int32)[0]] = ret['results']
            Bar.suffix = 'eval: [{0}/{1}]|Tot: {total:} |ETA: {eta:} '.format(
                ind, num_iters, total=bar.elapsed_td, eta=bar.eta_td)
            # for t in avg_time_stats:
            #     avg_time_stats[t].update(ret[t])
            #     Bar.suffix = Bar.suffix + '|{} {tm.val:.3f}s ({tm.avg:.3f}s) '.format(
            #         t, tm=avg_time_stats[t])
            bar.next()
        bar.finish()
        metrics, cls_eval = self.dataset.run_eval(results, self.opt.save_dir)
        return metrics, cls_eval
