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

import _init_paths
import argparse
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '3'
import torch
import json
import time
import cv2
import math

from sklearn import metrics
from scipy import interpolate
import numpy as np
from torchvision.transforms import transforms as T
import torch.nn.functional as F
import pickle

with open('tmp.pkl', 'rb') as pf:
    [gf, gpids] = pickle.load(pf)

embedding = gf
id_labels = torch.from_numpy(gpids)
n = id_labels.shape[0]
print(n, len(embedding))
assert len(embedding) == n

pdist = np.dot(embedding,embedding.T)
#pdist = torch.mm(embedding, embedding.t()).cpu().numpy()
gt = id_labels.expand(n, n).eq(id_labels.expand(n, n).t()).numpy()

up_triangle = np.where(np.triu(pdist) - np.eye(n) * pdist != 0)
pdist = pdist[up_triangle]
gt = gt[up_triangle]

far_levels = [1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1]
far, tar, threshold = metrics.roc_curve(gt, pdist)
interp = interpolate.interp1d(far, tar)
tar_at_far = [interp(x) for x in far_levels]
for f, fa in enumerate(far_levels):
    print('TPR@FAR={:.7f}: {:.4f}'.format(fa, tar_at_far[f]))
