from collections import OrderedDict
import time

import torch
import numpy as np
import pdb
from .evaluation import cmc, mean_ap, map_cmc
from .utils.meters import AverageMeter

from .utils import to_torch


def extract_cnn_feature(model, inputs, output_feature=None):
    model.eval()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    with torch.no_grad():
        inputs = to_torch(inputs)
        inputs = inputs.to(device)
        outputs = model(inputs, output_feature)
        outputs = outputs
    return outputs

def extract_features(model, data_loader, print_freq=1, output_feature=None):
    model.eval()
    batch_time = AverageMeter()
    data_time = AverageMeter()

    features = OrderedDict()
    labels = OrderedDict()

    end = time.time()
    for i, (imgs, fnames, pids, _) in enumerate(data_loader):
        data_time.update(time.time() - end)
        outputs = extract_cnn_feature(model, imgs, output_feature)
        for fname, output, pid in zip(fnames, outputs, pids):
            #print(fname)
            #if fname=='0005_c6_f0030883.jpg':
            #    pdb.set_trace()
            features[fname] = output
            labels[fname] = pid

        batch_time.update(time.time() - end)
        end = time.time()

    return features, labels

def pairwise_distance(query_features, gallery_features, query=None, gallery=None):
    x = torch.cat([query_features[f].unsqueeze(0) for f, _, _, _ in query], 0)
    y = torch.cat([gallery_features[f].unsqueeze(0) for f, _, _, _ in gallery], 0)

    m, n = x.size(0), y.size(0)
    x = x.view(m, -1)
    y = y.view(n, -1)

    dist = torch.pow(x, 2).sum(dim=1, keepdim=True).expand(m, n) + \
           torch.pow(y, 2).sum(dim=1, keepdim=True).expand(n, m).t()
    dist.addmm_(1, -2, x, y.t())
    # We use clamp to keep numerical stability
    _shape = np.shape(dist)
    if _shape[0]*_shape[1]>3000000:
        dist = torch.clamp(dist.to('cpu'), 1e-8, np.inf)
    else:
        dist = torch.clamp(dist, 1e-8, np.inf)
    return dist


def evaluate_all(distmat, query=None, gallery=None,
                 query_ids=None, gallery_ids=None,
                 query_cams=None, gallery_cams=None,wr=None,ep=None,suffix=None,
                 cmc_topk=(1, 5, 10, 20)):
    if query is not None and gallery is not None:
        query_ids = [pid for _, pid, _, _ in query]
        gallery_ids = [pid for _, pid, _, _ in gallery]
        query_cams = [cam for _, _, cam, _ in query]
        gallery_cams = [cam for _, _, cam, _ in gallery]
    else:
        assert (query_ids is not None and gallery_ids is not None
                and query_cams is not None and gallery_cams is not None)

    # Evaluation
    mAP, all_cmc = map_cmc(distmat, query_ids, gallery_ids, query_cams, gallery_cams)
    print('Mean AP: {:4.1%}'.format(mAP))
    if wr!=None:
        if suffix==None:
            wr.add_scalar("Eval/MeanAP", mAP,ep)
        else:
            wr.add_scalar("Eval/MeanAP-%s"%(suffix), mAP,ep)
    print('CMC Scores')
    for k in cmc_topk:
        print('  top-{:<4}{:12.1%}'
              .format(k, all_cmc[k - 1]))
        if wr!=None:
            if suffix==None:
                wr.add_scalar("CMC/Top-%d"%(k),all_cmc[k - 1],ep)
            else:
                wr.add_scalar("CMC/Top-%d-%s"%(k,suffix),all_cmc[k - 1],ep)
    return all_cmc[0]

class Evaluator(object):
    def __init__(self, model):
        super(Evaluator, self).__init__()
        self.model = model

    def evaluate(self, query_loader, gallery_loader, query, gallery, writter=None,epoch=None, output_feature=None,suffix=None):
        print('Extract feature from query images')
        query_features, _ = extract_features(self.model, query_loader, 1, output_feature)
        print('Extract feature from gallery images')
        gallery_features, _ = extract_features(self.model, gallery_loader, 1, output_feature)
        print('Computing distance between the two feature groups')
        distmat = pairwise_distance(query_features, gallery_features, query, gallery)
        print('Start to compute metric')
        return evaluate_all(distmat, query=query, gallery=gallery,wr=writter,ep=epoch,suffix=suffix)
