from __future__ import division, print_function, absolute_import
import time

import cv2
import numpy as np
import os.path as osp
import datetime
from collections import OrderedDict
import torch
from torch.nn import functional as F
from torch.utils.tensorboard import SummaryWriter

from torchreid import metrics
from torchreid.utils import (
    MetricMeter, AverageMeter, re_ranking, open_all_layers, save_checkpoint,
    open_specified_layers, visualize_ranked_results
)
from torchreid.losses import DeepSupervision


class Engine(object):
    r"""A generic base Engine class for both image- and video-reid.

    Args:
        datamanager (DataManager): an instance of ``torchreid.data.ImageDataManager``
            or ``torchreid.data.VideoDataManager``.
        use_gpu (bool, optional): use gpu. Default is True.
    """

    def __init__(self, datamanager, use_gpu=True):
        self.datamanager = datamanager
        self.train_loader = self.datamanager.train_loader
        self.test_loader = self.datamanager.test_loader
        self.use_gpu = (torch.cuda.is_available() and use_gpu)
        self.writer = None
        self.epoch = 0

        self.model = None
        self.optimizer = None
        self.scheduler = None

        self._models = OrderedDict()
        self._optims = OrderedDict()
        self._scheds = OrderedDict()

    def register_model(self, name='model', model=None, optim=None, sched=None):
        if self.__dict__.get('_models') is None:
            raise AttributeError(
                'Cannot assign model before super().__init__() call'
            )

        if self.__dict__.get('_optims') is None:
            raise AttributeError(
                'Cannot assign optim before super().__init__() call'
            )

        if self.__dict__.get('_scheds') is None:
            raise AttributeError(
                'Cannot assign sched before super().__init__() call'
            )

        self._models[name] = model
        self._optims[name] = optim
        self._scheds[name] = sched

    def get_model_names(self, names=None):
        names_real = list(self._models.keys())
        if names is not None:
            if not isinstance(names, list):
                names = [names]
            for name in names:
                assert name in names_real
            return names
        else:
            return names_real

    def save_model(self, epoch, rank1, save_dir, stime=None, is_best=False):
        names = self.get_model_names()

        for name in names:
            save_checkpoint(
                {
                    'state_dict': self._models[name].state_dict(),
                    'epoch': epoch + 1,
                    'rank1': rank1,
                    'optimizer': self._optims[name].state_dict(),
                    'scheduler': self._scheds[name].state_dict()
                },
                osp.join(save_dir, stime+name),
                # osp.join(save_dir, name),

                is_best=is_best
            )

    def set_model_mode(self, mode='train', names=None):
        assert mode in ['train', 'eval', 'test']
        names = self.get_model_names(names)

        for name in names:
            if mode == 'train':
                self._models[name].train()
            else:
                self._models[name].eval()

    def get_current_lr(self, names=None):
        names = self.get_model_names(names)
        name = names[0]
        return self._optims[name].param_groups[-1]['lr']

    def update_lr(self, names=None):
        names = self.get_model_names(names)

        for name in names:
            if self._scheds[name] is not None:
                self._scheds[name].step()

    def run(
        self,
        save_dir='log',
        max_epoch=0,
        start_epoch=0,
        print_freq=10,
        fixbase_epoch=0,
        open_layers=None,
        start_eval=0,
        eval_freq=-1,
        test_only= False,
        dist_metric='euclidean',
        normalize_feature=False,
        visrank=False,
        visrank_topk=10,
        use_metric_cuhk03=False,
        ranks=[1, 5, 10, 20],
        rerank=False,
        save_best_only=False,
        stime=time.strftime('%Y-%m-%d-%H-%M-%S'),
        compute_weight = False,
        save_distance = False
    ):
        r"""A unified pipeline for training and evaluating a model.

        Args:
            save_dir (str): directory to save model.
            max_epoch (int): maximum epoch.
            start_epoch (int, optional): starting epoch. Default is 0.
            print_freq (int, optional): print_frequency. Default is 10.
            fixbase_epoch (int, optional): number of epochs to train ``open_layers`` (new layers)
                while keeping base layers frozen. Default is 0. ``fixbase_epoch`` is counted
                in ``max_epoch``.
            open_layers (str or list, optional): layers (attribute names) open for training.
            start_eval (int, optional): from which epoch to start evaluation. Default is 0.
            eval_freq (int, optional): evaluation frequency. Default is -1 (meaning evaluation
                is only performed at the end of training).
            test_only (bool, optional): if True, only runs evaluation on test datasets.
                Default is False.
            dist_metric (str, optional): distance metric used to compute distance matrix
                between query and gallery. Default is "euclidean".
            normalize_feature (bool, optional): performs L2 normalization on feature vectors before
                computing feature distance. Default is False.
            visrank (bool, optional): visualizes ranked results. Default is False. It is recommended to
                enable ``visrank`` when ``test_only`` is True. The ranked images will be saved to
                "save_dir/visrank_dataset", e.g. "save_dir/visrank_market1501".
            visrank_topk (int, optional): top-k ranked images to be visualized. Default is 10.
            use_metric_cuhk03 (bool, optional): use single-gallery-shot setting for cuhk03.
                Default is False. This should be enabled when using cuhk03 classic split.
            ranks (list, optional): cmc ranks to be computed. Default is [1, 5, 10, 20].
            rerank (bool, optional): uses person re-ranking (by Zhong et al. CVPR'17).
                Default is False. This is only enabled when test_only=True.
            save_best_only (bool, optional): during training, save the best model on test set and last epoch.
                Default is True to save storage.
        """
        if visrank and not test_only:
            raise ValueError(
                'visrank can be set to True only if test_only=True'
            )

        if test_only:
            self.test(
                dist_metric=dist_metric,
                normalize_feature=normalize_feature,
                visrank=visrank,
                visrank_topk=visrank_topk,
                save_dir=save_dir,
                use_metric_cuhk03=use_metric_cuhk03,
                ranks=ranks,
                rerank=rerank,
                compute_weight=compute_weight,
                save_distance=save_distance
            )
            return

        if self.writer is None:
            self.writer = SummaryWriter(log_dir=save_dir)

        time_start = time.time()
        self.start_epoch = start_epoch
        self.max_epoch = max_epoch

        print('=> Start training')

        rank1_best = 0

        for self.epoch in range(self.start_epoch, self.max_epoch):
            print(self.epoch)
            self.train(
                print_freq=print_freq,
                fixbase_epoch=fixbase_epoch,
                open_layers=open_layers,
                save_dir=save_dir,
                stime=stime
            )

            if (self.epoch + 1) >= start_eval \
               and eval_freq > 0 \
               and (self.epoch+1) % eval_freq == 0 \
               and (self.epoch + 1) != self.max_epoch:
                rank1 = self.test(
                    dist_metric=dist_metric,
                    normalize_feature=normalize_feature,
                    visrank=visrank,
                    visrank_topk=visrank_topk,
                    save_dir=save_dir,
                    use_metric_cuhk03=use_metric_cuhk03,
                    ranks=ranks
                )
                # self.save_model(self.epoch, rank1, save_dir, stime)

                if not save_best_only:
                    self.save_model(self.epoch, rank1, save_dir, stime, is_best=True)
                    # save best epoch on test set
                elif rank1 >= rank1_best:
                    rank1_best = rank1
                    self.save_model(self.epoch, rank1, save_dir, stime, is_best=True)

        if self.max_epoch > 0:
            print('=> Final test')
            rank1 = self.test(
                dist_metric=dist_metric,
                normalize_feature=normalize_feature,
                visrank=visrank,
                visrank_topk=visrank_topk,
                save_dir=save_dir,
                use_metric_cuhk03=use_metric_cuhk03,
                ranks=ranks
            )
            self.save_model(self.epoch, rank1, save_dir, stime)

        elapsed = round(time.time() - time_start)
        elapsed = str(datetime.timedelta(seconds=elapsed))
        print('Elapsed {}'.format(elapsed))
        if self.writer is not None:
            self.writer.close()

    def train(self, print_freq=10, fixbase_epoch=0, open_layers=None, save_dir='/mnt/data/code/deep-person-reid-master/log_test2', stime=None):
        losses = MetricMeter()
        batch_time = AverageMeter()
        data_time = AverageMeter()

        self.set_model_mode('train')

        self.two_stepped_transfer_learning(
            self.epoch, fixbase_epoch, open_layers
        )

        self.num_batches = len(self.train_loader)
        end = time.time()
        for self.batch_idx, data in enumerate(self.train_loader):
            data_time.update(time.time() - end)
            loss_summary = self.forward_backward(data)
            batch_time.update(time.time() - end)
            losses.update(loss_summary)

            if (self.batch_idx + 1) % print_freq == 0:
                nb_this_epoch = self.num_batches - (self.batch_idx + 1)
                nb_future_epochs = (
                    self.max_epoch - (self.epoch + 1)
                ) * self.num_batches
                eta_seconds = batch_time.avg * (nb_this_epoch+nb_future_epochs)
                eta_str = str(datetime.timedelta(seconds=int(eta_seconds)))
                print(
                    'epoch: [{0}/{1}][{2}/{3}]\t'
                    'time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                    'data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                    'eta {eta}\t'
                    '{losses}\t'
                    'lr {lr:.6f}'.format(
                        self.epoch + 1,
                        self.max_epoch,
                        self.batch_idx + 1,
                        self.num_batches,
                        batch_time=batch_time,
                        data_time=data_time,
                        eta=eta_str,
                        losses=losses,
                        lr=self.get_current_lr()
                    )
                )

            if self.writer is not None:
                n_iter = self.epoch * self.num_batches + self.batch_idx
                self.writer.add_scalar('Train/time', batch_time.avg, n_iter)
                self.writer.add_scalar('Train/data', data_time.avg, n_iter)
                for name, meter in losses.meters.items():
                    self.writer.add_scalar('Train/' + name, meter.avg, n_iter)
                self.writer.add_scalar(
                    'Train/lr', self.get_current_lr(), n_iter
                )

            end = time.time()

        self.update_lr()
        # if(self.epoch >= 200):
        #     rank1 = self.test()
        #     self.save_model(self.epoch, rank1, save_dir, stime)

    def forward_backward(self, data):
        raise NotImplementedError

    def test(
        self,
        dist_metric='euclidean',
        normalize_feature=False,
        visrank=False,
        visrank_topk=10,
        save_dir='',
        use_metric_cuhk03=False,
        ranks=[1, 5, 10, 20],
        rerank=False,
        compute_weight=False,
        save_distance=False
    ):
        r"""Tests model on target datasets.

        .. note::

            This function has been called in ``run()``.

        .. note::

            The test pipeline implemented in this function suits both image- and
            video-reid. In general, a subclass of Engine only needs to re-implement
            ``extract_features()`` and ``parse_data_for_eval()`` (most of the time),
            but not a must. Please refer to the source code for more details.
        """
        self.set_model_mode('eval')
        targets = list(self.test_loader.keys())

        for name in targets:
            domain = 'source' if name in self.datamanager.sources else 'target'
            print('##### Evaluating {} ({}) #####'.format(name, domain))
            query_loader = self.test_loader[name]['query']
            gallery_loader = self.test_loader[name]['gallery']
            rank1, mAP = self._evaluate(
                dataset_name=name,
                query_loader=query_loader,
                gallery_loader=gallery_loader,
                dist_metric=dist_metric,
                normalize_feature=normalize_feature,
                visrank=visrank,
                visrank_topk=visrank_topk,
                save_dir=save_dir,
                use_metric_cuhk03=use_metric_cuhk03,
                ranks=ranks,
                rerank=rerank,
                compute_weight=compute_weight,
                save_distance=save_distance
            )

            if self.writer is not None:
                self.writer.add_scalar(f'Test/{name}/rank1', rank1, self.epoch)
                self.writer.add_scalar(f'Test/{name}/mAP', mAP, self.epoch)

        return rank1

    @torch.no_grad()
    def _evaluate(
        self,
        dataset_name='',
        query_loader=None,
        gallery_loader=None,
        dist_metric='euclidean',
        normalize_feature=False,
        visrank=False,
        visrank_topk=10,
        save_dir='',
        use_metric_cuhk03=False,
        ranks=[1, 5, 10, 20],
        rerank=False,
        compute_weight=False,
        save_distance=False
    ):
        batch_time = AverageMeter()

        def _feature_extraction(data_loader):
            f_, pids_, camids_ = [], [], []

            #获取对应尺度
            size1_, size2_, size_ = [], [], []
            imgpaths_ = []

            for batch_idx, data in enumerate(data_loader):
                # imgs, pids, camids = self.parse_data_for_eval(data)
                imgs, pids, camids, imgpaths = self.parse_data_for_eval(data)

                if self.use_gpu:
                    imgs = imgs.cuda()
                end = time.time()

                # for i in range(len(imgpaths)):
                #     img = cv2.imread(imgpaths[i])
                #     size1 = img.shape[1]
                #     size2 = img.shape[0]
                #     size = size1*size2
                #     size_.append(size)
                #     size1_.append(size1)
                #     size2_.append(size2)
                # imgpaths_.extend(imgpaths)

                features = self.extract_features(imgs)
                batch_time.update(time.time() - end)
                features = features.cpu().clone()
                f_.append(features)
                pids_.extend(pids)
                camids_.extend(camids)
            f_ = torch.cat(f_, 0)
            pids_ = np.asarray(pids_)
            camids_ = np.asarray(camids_)
            return f_, pids_, camids_
            # return f_, pids_, camids_, size1_, size2_, size_, imgpaths_


        print('Extracting features from query set ...')
        qf, q_pids, q_camids, = _feature_extraction(query_loader)
        # qf, q_pids, q_camids, qsize1_, qsize2_, qsize_, qimgpaths = _feature_extraction(query_loader)
        print('Done, obtained {}-by-{} matrix'.format(qf.size(0), qf.size(1)))

        print('Extracting features from gallery set ...')
        gf, g_pids, g_camids, = _feature_extraction(gallery_loader)
        # gf, g_pids, g_camids, gsize1_, gsize2_, gsize_, gimgpaths = _feature_extraction(gallery_loader)
        print('Done, obtained {}-by-{} matrix'.format(gf.size(0), gf.size(1)))

        print('Speed: {:.4f} sec/batch'.format(batch_time.avg))

        if normalize_feature:
            print('Normalzing features with L2 norm ...')
            qf = F.normalize(qf, p=2, dim=1)
            gf = F.normalize(gf, p=2, dim=1)

        print(
            'Computing distance matrix with metric={} ...'.format(dist_metric)
        )
        distmat = metrics.compute_distance_matrix(qf, gf, dist_metric)
        distmat = distmat.numpy()

        # compute_weight = self.comput_weight
        # save_distance = self.save_distance
        # f1 = True #加权距离向量
        # f2 = True #输出距离向量
        # f1 = False
        # f2 = False
        # if compute_weight:
        #     with open(osp.join('/mnt/data/code/reidlog/dist', 'imgpaths.csv'), 'w') as f:
        #         f.write(str(gimgpaths))
        #
        #     gsize_ = np.array(gsize_)
        #     qsize1_ = np.array(qsize1_)
        #     qsize2_ = np.array(qsize2_)
        #
        #     qsize_ = np.array(qsize_)
        #
        #     # 平均尺度
        #     average = np.average(gsize_)
        #     bili = qsize_/average
        #     r = np.sqrt(qsize_ / average)
        #     r_m1 = r - 1
        #     r_m2 = r - 1/2
        #     r_m3 = r - 1/3
        #     r_m4 = r - 1/4
        #
        #     # # --------更改公式
        #     # dt = 0.9
        #     # print('更改公式后计算权重，dt='+ str(dt))
        #     # r_m1 = np.maximum(r_m1, -r_m1)
        #     # r_m2 = np.maximum(r_m2, -r_m2)
        #     # r_m3 = np.maximum(r_m3, -r_m3)
        #     # r_m4 = np.maximum(r_m4, -r_m4)
        #     #
        #     # w1 = np.exp(-pow(dt, 2) * r_m1)
        #     # w2 = np.exp(-pow(dt, 2) * r_m2)
        #     # w3 = np.exp(-pow(dt, 2) * r_m3)
        #     # w4 = np.exp(-pow(dt, 2) * r_m4)
        #     # # --------更改公式
        #
        #     dt = 0.2
        #     print('计算权重，dt=' + str(dt))
        #     w1 = np.exp(-pow(dt,2)*pow(r_m1,2))
        #     w2 = np.exp(-pow(dt,2)*pow(r_m2,2))
        #     w3 = np.exp(-pow(dt,2)*pow(r_m3,2))
        #     w4 = np.exp(-pow(dt,2)*pow(r_m4,2))
        #
        #     #依次为 查询图像宽、长、像素点数、qeury与gallery像素点平均的比值、对应r
        #     # r_m2: r-m2
        #     size = np.vstack((qsize1_, qsize2_, qsize_, bili, r, r_m1, r_m2, r_m3, r_m4, w1, w2, w3, w4))
        #     # size = np.vstack((qsize1_, qsize2_, qsize_, bili, r, r_m2, r_m3, r_m4, w2, w3, w4))
        #     trans = np.transpose(size)
        #     # size = np.vstack((qsize2_))
        #     # np.savetxt(osp.join('/mnt/data/code/reidlog/dist', 'gsize.csv'), size, delimiter=',')
        #     # np.savetxt(osp.join('/mnt/data/code/reidlog/dist', 'transposegsize.csv'), trans, delimiter=',')
        #     np.savetxt(osp.join(save_dir, 'transposegsize.csv'), trans, delimiter=',')
        #     w = np.vstack((w1, w2, w3, w4))
        #
        #
        #     # dist2 = np.loadtxt(open('/mnt/data/code/reidlog/dist/mlr_m/mnt/data/code/reidlog/MTA_reid/x2arket1501_split0/distx2.csv', "rb"), delimiter=",", skiprows=0)
        #     # dist3 = np.loadtxt(open('/mnt/data/code/reidlog/dist/mlr_market1501_split0/distx3.csv', "rb"), delimiter=",", skiprows=0)
        #     # dist4 = np.loadtxt(open('/mnt/data/code/reidlog/dist/mlr_market1501_split0/distx4.csv', "rb"), delimiter=",", skiprows=0)
        #
        #     # dist0 = distmat
        #     # dist1 = np.loadtxt('/mnt/data/code/reidlog/mlr_market1501_hr_dsx2/hr_dsx2/dist.txt')
        #     # dist2 = np.loadtxt('/mnt/data/code/reidlog/mlr_market1501_hr_dsx2/x2sr/dist.txt')
        #     # dist3 = np.loadtxt('/mnt/data/code/reidlog/mlr_market1501_hr_dsx2/x3sr/dist.txt')
        #     # dist4 = np.loadtxt('/mnt/data/code/reidlog/mlr_market1501_hr_dsx2/x4sr/dist.txt')
        #     print('加载距离向量......')
        #     # dist1 = np.loadtxt(('/mnt/data/code/reidlog/MTA-reid/dist.txt'))
        #     dist2 = np.loadtxt('/mnt/data/code/reidlog/MTA_reid/x2/dist.txt')
        #     # dist3 = np.loadtxt('/mnt/data/code/reidlog/MTA_reid/x3/dist.txt')
        #     # dist4 = np.loadtxt('/mnt/data/code/reidlog/MTA_reid/x4/dist.txt')
        #
        #     #------距离向量加权-------
        #     print('距离向量加权......')
        #     # for i in range(dist1.shape[1]):
        #     #     dist1[:, i] = np.multiply(dist1[:, i], w2).reshape(dist1.shape[0],)
        #
        #     for i in range(dist2.shape[1]):
        #         dist2[:, i] = np.multiply(dist2[:, i], w2).reshape(dist2.shape[0],)
        #     # for i in range(dist3.shape[1]):
        #     #     dist3[:, i] = np.multiply(dist3[:, i], w3).reshape(dist3.shape[0],)
        #     # for i in range(dist4.shape[1]):
        #     #     dist4[:, i] = np.multiply(dist4[:, i], w4).reshape(dist4.shape[0],)
        #     # ------距离向量加权-------
        #
        #
        #     # print('distx2')
        #     # distmat = dist2
        #     # print('distx234')
        #     # distmat = dist2 + dist3 + dist4
        #     # print('distx1234')
        #     # distmat = dist2 + dist3 + dist4 + dist1
        #     # print('dist0+x1234')
        #     # distmat = dist2 + dist3 + dist4 + dist1 + dist0
        #
        # if save_distance:
        #     print('Saving distance!')
        #     #输出距离向量
        #     distpath = osp.join(save_dir, 'dist.txt')
        #     # np.savetxt(osp.join('/mnt/data/code/reidlog/dist', 'dist.csv'), distmat, delimiter=',')
        #     np.savetxt(distpath, distmat)
        #     # np.savetxt(osp.join(save_dir, 'w.txt'), w)
        #     # np.savetxt(osp.join(save_dir, 'transposegsize.csv'), trans)
        #     print('Save success!')


        if rerank:
            print('Applying person re-ranking ...')
            distmat_qq = metrics.compute_distance_matrix(qf, qf, dist_metric)
            distmat_gg = metrics.compute_distance_matrix(gf, gf, dist_metric)
            distmat = re_ranking(distmat, distmat_qq, distmat_gg)

        print('Computing CMC and mAP ...')
        cmc, mAP = metrics.evaluate_rank(
            distmat,
            q_pids,
            g_pids,
            q_camids,
            g_camids,
            use_metric_cuhk03=use_metric_cuhk03
        )


        print('** Results **')
        print('mAP: {:.1%}'.format(mAP))
        print('CMC curve')
        for r in ranks:
            print('Rank-{:<3}: {:.1%}'.format(r, cmc[r - 1]))

        if visrank:
            visualize_ranked_results(
                distmat,
                self.datamanager.fetch_test_loaders(dataset_name),
                self.datamanager.data_type,
                width=self.datamanager.width,
                height=self.datamanager.height,
                save_dir=osp.join(save_dir, 'visrank_' + dataset_name),
                topk=visrank_topk
            )

        return cmc[0], mAP

    def compute_loss(self, criterion, outputs, targets):
        if isinstance(outputs, (tuple, list)):
            loss = DeepSupervision(criterion, outputs, targets)
        else:
            loss = criterion(outputs, targets)
        return loss

    def extract_features(self, input):
        return self.model(input)

    def parse_data_for_train(self, data):
        imgs = data['img']
        pids = data['pid']
        return imgs, pids

    def parse_data_for_eval(self, data):
        imgs = data['img']
        pids = data['pid']
        camids = data['camid']

        imgpath = data['impath']


        return imgs, pids, camids,imgpath

    def two_stepped_transfer_learning(
        self, epoch, fixbase_epoch, open_layers, model=None
    ):
        """Two-stepped transfer learning.

        The idea is to freeze base layers for a certain number of epochs
        and then open all layers for training.

        Reference: https://arxiv.org/abs/1611.05244
        """
        model = self.model if model is None else model
        if model is None:
            return

        if (epoch + 1) <= fixbase_epoch and open_layers is not None:
            print(
                '* Only train {} (epoch: {}/{})'.format(
                    open_layers, epoch + 1, fixbase_epoch
                )
            )
            open_specified_layers(model, open_layers)
        else:
            open_all_layers(model)
