# Copyright 2021 Huawei Technologies Co., Ltd

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

#     http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import argparse

import os

import sys

import json

import math

import numpy as np

from tqdm import tqdm

from scipy import interpolate

from sklearn.model_selection import KFold





def read_pairs(pairs_filename):

    pairs = []

    with open(pairs_filename, 'r') as f:

        for line in f.readlines()[1:]:

            pair = line.strip().split()

            pairs.append(pair)

    return np.array(pairs, dtype=object)





def load_json(json_path):

    with open(json_path) as f:

        return json.load(f)





def add_extension(path):

    if os.path.exists(path+'.jpg'):

        return path+'.jpg'

    elif os.path.exists(path+'.png'):

        return path+'.png'

    else:

        raise RuntimeError('No file "%s" with extension png or jpg.' % path)





def get_paths(lfw_dir, pairs):

    nrof_skipped_pairs = 0

    path_list = []

    issame_list = []

    for pair in pairs:

        if len(pair) == 3:

            path0 = add_extension(os.path.join(lfw_dir, pair[0], pair[0] + '_' + '%04d' % int(pair[1])))

            path1 = add_extension(os.path.join(lfw_dir, pair[0], pair[0] + '_' + '%04d' % int(pair[2])))

            issame = True

        elif len(pair) == 4:

            path0 = add_extension(os.path.join(lfw_dir, pair[0], pair[0] + '_' + '%04d' % int(pair[1])))

            path1 = add_extension(os.path.join(lfw_dir, pair[2], pair[2] + '_' + '%04d' % int(pair[3])))

            issame = False

        if os.path.exists(path0) and os.path.exists(path1):    # Only add the pair if both paths exist

            path_list += (path0, path1)

            issame_list.append(issame)

        else:

            nrof_skipped_pairs += 1

    if nrof_skipped_pairs > 0:

        print('Skipped %d image pairs' % nrof_skipped_pairs)



    return path_list, issame_list





def face_postprocess(crop_paths, result_dir):

    num_bins = len(os.listdir(result_dir))

    embeddings = []

    flag_file = os.path.join(result_dir, "{}_output_0.bin".format(0))

    for idx in tqdm(range(num_bins)):

        if not os.path.exists(flag_file):

            xb_path = os.path.join(result_dir, "{}_1.bin".format(idx))

        else:

            xb_path = os.path.join(result_dir, "{}_output_0.bin".format(idx))

        xb_data = np.fromfile(xb_path, dtype=np.float32).reshape(-1, 512)

        embeddings.extend(xb_data)



    embeddings_dict = dict(zip(crop_paths, embeddings))

    return embeddings_dict





def evaluate(embeddings, actual_issame, nrof_folds=10, distance_metric=0, subtract_mean=False):

    # Calculate evaluation metrics

    thresholds = np.arange(0, 4, 0.01)

    embeddings1 = embeddings[0::2]

    embeddings2 = embeddings[1::2]

    tpr, fpr, accuracy, fp, fn = calculate_roc(thresholds, embeddings1, embeddings2,

                                                np.asarray(actual_issame), nrof_folds=nrof_folds, distance_metric=distance_metric, subtract_mean=subtract_mean)

    thresholds = np.arange(0, 4, 0.001)

    val, val_std, far = calculate_val(thresholds, embeddings1, embeddings2,

                                      np.asarray(actual_issame), 1e-3, nrof_folds=nrof_folds, distance_metric=distance_metric, subtract_mean=subtract_mean)

    return tpr, fpr, accuracy, val, val_std, far, fp, fn





def calculate_roc(thresholds, embeddings1, embeddings2, actual_issame, nrof_folds=10, distance_metric=0, subtract_mean=False):

    assert(embeddings1.shape[0] == embeddings2.shape[0])

    assert(embeddings1.shape[1] == embeddings2.shape[1])

    nrof_pairs = min(len(actual_issame), embeddings1.shape[0])

    nrof_thresholds = len(thresholds)

    k_fold = KFold(n_splits=nrof_folds, shuffle=False)



    tprs = np.zeros((nrof_folds,nrof_thresholds))

    fprs = np.zeros((nrof_folds,nrof_thresholds))

    accuracy = np.zeros((nrof_folds))



    is_false_positive = []

    is_false_negative = []



    indices = np.arange(nrof_pairs)



    for fold_idx, (train_set, test_set) in enumerate(k_fold.split(indices)):

        if subtract_mean:

            mean = np.mean(np.concatenate([embeddings1[train_set], embeddings2[train_set]]), axis=0)

        else:

            mean = 0.0

            dist = distance(embeddings1-mean, embeddings2-mean, distance_metric)



        # Find the best threshold for the fold

        acc_train = np.zeros((nrof_thresholds))

        for threshold_idx, threshold in enumerate(thresholds):

            _, _, acc_train[threshold_idx], _, _ = calculate_accuracy(threshold, dist[train_set], actual_issame[train_set])

            best_threshold_index = np.argmax(acc_train)

        for threshold_idx, threshold in enumerate(thresholds):

            tprs[fold_idx, threshold_idx], fprs[fold_idx, threshold_idx], _, _, _ = calculate_accuracy(threshold, dist[test_set], actual_issame[test_set])

            _, _, accuracy[fold_idx], is_fp, is_fn = calculate_accuracy(thresholds[best_threshold_index], dist[test_set], actual_issame[test_set])



        tpr = np.mean(tprs, 0)

        fpr = np.mean(fprs, 0)

        is_false_positive.extend(is_fp)

        is_false_negative.extend(is_fn)



    return tpr, fpr, accuracy, is_false_positive, is_false_negative





def calculate_accuracy(threshold, dist, actual_issame):

    predict_issame = np.less(dist, threshold)

    tp = np.sum(np.logical_and(predict_issame, actual_issame))

    fp = np.sum(np.logical_and(predict_issame, np.logical_not(actual_issame)))

    tn = np.sum(np.logical_and(np.logical_not(predict_issame), np.logical_not(actual_issame)))

    fn = np.sum(np.logical_and(np.logical_not(predict_issame), actual_issame))



    is_fp = np.logical_and(predict_issame, np.logical_not(actual_issame))

    is_fn = np.logical_and(np.logical_not(predict_issame), actual_issame)



    tpr = 0 if (tp + fn == 0) else float(tp) / float(tp+fn)

    fpr = 0 if (fp + tn == 0) else float(fp) / float(fp+tn)

    acc = float(tp+tn)/dist.size

    return tpr, fpr, acc, is_fp, is_fn





def calculate_val(thresholds, embeddings1, embeddings2, actual_issame, far_target, nrof_folds=10, distance_metric=0, subtract_mean=False):

    assert(embeddings1.shape[0] == embeddings2.shape[0])

    assert(embeddings1.shape[1] == embeddings2.shape[1])

    nrof_pairs = min(len(actual_issame), embeddings1.shape[0])

    nrof_thresholds = len(thresholds)

    k_fold = KFold(n_splits=nrof_folds, shuffle=False)



    val = np.zeros(nrof_folds)

    far = np.zeros(nrof_folds)



    indices = np.arange(nrof_pairs)



    for fold_idx, (train_set, test_set) in enumerate(k_fold.split(indices)):

        if subtract_mean:

            mean = np.mean(np.concatenate([embeddings1[train_set], embeddings2[train_set]]), axis=0)

        else:

            mean = 0.0

            dist = distance(embeddings1-mean, embeddings2-mean, distance_metric)



        # Find the threshold that gives FAR = far_target

        far_train = np.zeros(nrof_thresholds)

        for threshold_idx, threshold in enumerate(thresholds):

            _, far_train[threshold_idx] = calculate_val_far(threshold, dist[train_set], actual_issame[train_set])

        if np.max(far_train)>=far_target:

            f = interpolate.interp1d(far_train, thresholds, kind='slinear')

            threshold = f(far_target)

        else:

            threshold = 0.0



        val[fold_idx], far[fold_idx] = calculate_val_far(threshold, dist[test_set], actual_issame[test_set])



    val_mean = np.mean(val)

    far_mean = np.mean(far)

    val_std = np.std(val)

    return val_mean, val_std, far_mean





def distance(embeddings1, embeddings2, distance_metric=0):

    if distance_metric==0:

        # Euclidian distance

        diff = np.subtract(embeddings1, embeddings2)

        dist = np.sum(np.square(diff),1)

    elif distance_metric==1:

        # Distance based on cosine similarity

        dot = np.sum(np.multiply(embeddings1, embeddings2), axis=1)

        norm = np.linalg.norm(embeddings1, axis=1) * np.linalg.norm(embeddings2, axis=1)

        similarity = dot / norm

        dist = np.arccos(similarity) / math.pi

    else:

        raise 'Undefined distance metric %d' % distance_metric



    return dist





def calculate_val_far(threshold, dist, actual_issame):

    predict_issame = np.less(dist, threshold)

    true_accept = np.sum(np.logical_and(predict_issame, actual_issame))

    false_accept = np.sum(np.logical_and(predict_issame, np.logical_not(actual_issame)))

    n_same = np.sum(actual_issame)

    n_diff = np.sum(np.logical_not(actual_issame))

    val = float(true_accept) / float(n_same)

    far = float(false_accept) / float(n_diff)

    return val, far





if __name__ == '__main__':

    parser = argparse.ArgumentParser()

    parser.add_argument('--pair_path', default='./data/pairs.txt', type=str, help='path for pair gt label')

    parser.add_argument('--crop_dir', type=str, help='cropped image save path')

    parser.add_argument('--test_dir', type=str, help='test file path')

    parser.add_argument('--ONet_output_dir', type=str, help='preprocess bin files save path')

    arg = parser.parse_args()

    embedding_output_path = arg.test_dir

    pairs = read_pairs(arg.pair_path)

    crop_paths = load_json(arg.ONet_output_dir)

    crop_dir = arg.crop_dir

    path_list, _ = get_paths(crop_dir, pairs)

    embeddings_dict = face_postprocess(crop_paths, embedding_output_path)

    if list(embeddings_dict.keys())[0][:2] == './':

        embeddings = np.array([embeddings_dict['./' + os.path.relpath(path)] for path in path_list])

    else:

        embeddings = np.array([embeddings_dict[os.path.relpath(path)] for path in path_list])

    path_list, issame_list = get_paths(crop_dir, pairs)

    tpr, fpr, accuracy, val, val_std, far, fp, fn = evaluate(embeddings, issame_list)

    print("accuracy:", accuracy)

    print("mean accuracy:", np.mean(accuracy))