import argparse
import numpy as np
from sklearn.model_selection import StratifiedKFold
import sklearn
import cv2
import datetime
import mxnet as mx
from mxnet import ndarray as nd
import pandas as pd
from numpy import linalg as line
import logging

logging.basicConfig(
    format="%(asctime)s %(message)s", datefmt="%m/%d/%Y %I:%M:%S %p", level=logging.INFO
)


class FaceVerification:
    def __init__(self, model=None, batch_size=32, data_dir=None):
        super().__init__()
        logging.info("Face Verification for RFW.")
        self.data_dir = data_dir
        self.image_size = 112
        self.batch_size = batch_size
        self.model = model

    def load_model(self, model_dir=None):
        logging.info("Model Loading")
        ctx = mx.gpu(0)
        sym, arg_params, aux_params = mx.model.load_checkpoint(model_dir, 1)
        all_layers = sym.get_internals()
        sym = all_layers["fc1_output"]
        self.model = mx.mod.Module(symbol=sym, context=ctx, label_names=None)
        self.model.bind(
            data_shapes=[
                ("data", (self.batch_size, 3, self.image_size, self.image_size))
            ]
        )
        self.model.set_params(arg_params, aux_params)
        return self.model

    def load_images(self, inp_csv_file):
        logging.info("Image Data Loading")
        issame_list, data_list = [], []
        pairs = pd.read_csv(inp_csv_file)
        # data_list = list(
        #     np.empty((2, pairs.shape[0] * 2, 3, self.image_size, self.image_size))
        # )
        for flip in [0, 1]:
            data = nd.empty((pairs.shape[0] * 2, 3, self.image_size, self.image_size))
            data_list.append(data)

        j = 0
        for i, row in pairs.iterrows():

            if i % 1000 == 0:
                logging.info("processing {}".format(i))

            issame_list.append(row.issame)
            path1 = "{}/{}/{}_{:04d}.jpg".format(
                self.data_dir,
                row.Class_ID_s1,
                row.Class_ID_s1.split("/")[1],
                int(row.img_id_s1),
            )
            path2 = "{}/{}/{}_{:04d}.jpg".format(
                self.data_dir,
                row.Class_ID_s2,
                row.Class_ID_s2.split("/")[1],
                int(row.img_id_s2),
            )
            im1 = cv2.imread(path1)
            im1 = cv2.cvtColor(im1, cv2.COLOR_BGR2RGB)
            im1 = np.transpose(im1, (2, 0, 1))  # 3*112*112, RGB
            im1 = mx.nd.array(im1)

            im2 = cv2.imread(path2)
            im2 = cv2.cvtColor(im2, cv2.COLOR_BGR2RGB)
            im2 = np.transpose(im2, (2, 0, 1))  # 3*112*112, RGB
            im2 = mx.nd.array(im2)

            for flip in [0, 1]:

                if flip == 1:
                    im1 = mx.ndarray.flip(im1, 2)
                data_list[flip][j][:] = im1

            for flip in [0, 1]:
                if flip == 1:
                    im2 = mx.ndarray.flip(im2, 2)
                data_list[flip][j + 1][:] = im2
                # data_list[flip][i][:] = img

            j = j + 2
        # bins shape should be 2,12000,3,112,112

        # data = np.asarray(data_list)
        self.issame = np.asarray(issame_list)

        self.data = data_list
        logging.info("Pairs are loaded, shape: 2x{}.".format(self.data[0].shape))
        return self.data, self.issame, pairs.shape

    def clean_data(self):
        self.data = None
        self.issame = None

    def verify(self, model=None):
        data_list = self.data

        embeddings_list = []
        time_consumed = 0
        _label = nd.ones((self.batch_size,))
        for i in range(len(data_list)):
            data = data_list[i]
            embeddings = None
            ba = 0
            while ba < data.shape[0]:
                bb = min(ba + self.batch_size, data.shape[0])
                count = bb - ba
                _data = nd.slice_axis(data, axis=0, begin=bb - self.batch_size, end=bb)
                time0 = datetime.datetime.now()

                db = mx.io.DataBatch(data=(_data,), label=(_label,))

                self.model.forward(db, is_train=False)
                net_out = self.model.get_outputs()
                _embeddings = net_out[0].asnumpy()
                time_now = datetime.datetime.now()
                diff = time_now - time0
                time_consumed += diff.total_seconds()
                if embeddings is None:
                    embeddings = np.zeros((data.shape[0], _embeddings.shape[1]))
                embeddings[ba:bb, :] = _embeddings[(self.batch_size - count) :, :]
                ba = bb
            embeddings_list.append(embeddings)

        _xnorm = 0.0
        _xnorm_cnt = 0
        for embed in embeddings_list:
            for i in range(embed.shape[0]):
                _em = embed[i]
                _norm = np.linalg.norm(_em)
                _xnorm += _norm
                _xnorm_cnt += 1
        _xnorm /= _xnorm_cnt

        acc1 = 0.0
        std1 = 0.0
        embeddings = embeddings_list[0] + embeddings_list[1]
        embeddings = sklearn.preprocessing.normalize(embeddings)
        print(embeddings.shape)
        print("infer time", time_consumed)

        tpr, fpr, accuracy, best_thresholds = self.evaluate(
            embeddings, self.issame, nrof_folds=10
        )

        acc2, std2 = np.mean(accuracy), np.std(accuracy)
        logging.info("Accuracy {}".format(acc2))
        return tpr, fpr, acc2, std2

    def evaluate(self, embeddings, actual_issame, nrof_folds=10):
        # Calculate evaluation metrics
        thresholds = np.arange(-1, 1, 0.001)
        embeddings1 = embeddings[0::2]
        embeddings2 = embeddings[1::2]
        tpr, fpr, accuracy, best_thresholds = self.calculate_roc(
            thresholds,
            embeddings1,
            embeddings2,
            np.asarray(actual_issame),
            nrof_folds=nrof_folds,
        )

        return tpr, fpr, accuracy, best_thresholds

    def calculate_roc(
        self, thresholds, embeddings1, embeddings2, actual_issame, nrof_folds=10
    ):

        assert embeddings1.shape[1] == embeddings2.shape[1]
        nrof_pairs = min(len(actual_issame), embeddings1.shape[0])
        nrof_thresholds = len(thresholds)
        # k_fold = LFold(n_splits=nrof_folds, shuffle=False)
        k_fold = StratifiedKFold(n_splits=nrof_folds, shuffle=False)

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

        tnrs = np.zeros((nrof_folds, nrof_thresholds))
        fnrs = np.zeros((nrof_folds, nrof_thresholds))

        f1s = np.zeros((nrof_folds))

        accuracy = np.zeros((nrof_folds))
        indices = np.arange(nrof_pairs)

        veclist = np.concatenate((embeddings1, embeddings2), axis=0)
        meana = np.mean(veclist, axis=0)
        embeddings1 -= meana
        embeddings2 -= meana
        dist = np.sum(embeddings1 * embeddings2, axis=1)
        dist = dist / line.norm(embeddings1, axis=1) / line.norm(embeddings2, axis=1)

        for fold_idx, (train_set, test_set) in enumerate(
            k_fold.split(indices, actual_issame)
        ):
            # print(train_set.shape, actual_issame[train_set].sum())
            # print(test_set.shape, actual_issame[test_set].sum())

            # Find the best threshold for the fold
            acc_train = np.zeros((nrof_thresholds))
            for threshold_idx, threshold in enumerate(thresholds):
                _, _, _, _, acc_train[threshold_idx], f1 = self.calculate_accuracy(
                    threshold, dist[train_set], actual_issame[train_set]
                )
            best_threshold_index = np.argmax(acc_train)
            # print('threshold', thresholds[best_threshold_index])
            for threshold_idx, threshold in enumerate(thresholds):
                (
                    tprs[fold_idx, threshold_idx],
                    fprs[fold_idx, threshold_idx],
                    tnrs[fold_idx, threshold_idx],
                    fnrs[fold_idx, threshold_idx],
                    _,
                    _,
                ) = self.calculate_accuracy(
                    threshold, dist[test_set], actual_issame[test_set]
                )

            _, _, _, _, accuracy[fold_idx], f1s[fold_idx] = self.calculate_accuracy(
                thresholds[best_threshold_index],
                dist[test_set],
                actual_issame[test_set],
            )
        tpr = np.mean(tprs, 0)[best_threshold_index]
        fpr = np.mean(fprs, 0)[best_threshold_index]

        # tnr = np.mean(tnrs, 0)[best_threshold_index]
        # fnr = np.mean(fnrs, 0)[best_threshold_index]

        return tpr, fpr, accuracy, thresholds[best_threshold_index]

    def calculate_accuracy(self, threshold, dist, actual_issame):
        predict_issame = np.less(dist, threshold)
        actual_issame = np.less(actual_issame, 0.5)

        tn, fp, fn, tp = sklearn.metrics.confusion_matrix(
            actual_issame, predict_issame
        ).ravel()

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

        tnr = 0 if (fp + tn == 0) else float(tn) / float(fp + tn)
        fnr = 0 if (fn + tp == 0) else float(fn) / float(fn + tp)

        acc = float(tp + tn) / dist.size
        f1 = sklearn.metrics.f1_score(predict_issame, actual_issame)

        return tpr, fpr, tnr, fnr, acc, f1


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Face Verification for RFW")
    parser.add_argument(
        "--data_dir", type=str, default="RFW/test/aligned_data", help="dataset root"
    )
    parser.add_argument(
        "--pair_file",
        type=str,
        default="./AttributePairs/eye_narrow_pairs_6000_selected.csv",
        help="pair file to test",
    )
    parser.add_argument(
        "--model_dir", type=str, default="/model/", help="pre-trained model directory"
    )

    parser.add_argument("--batch_size", type=int, default="32", help="batch_size")
    args = parser.parse_args()

    validation = FaceVerification(
        batch_size=args.batch_size, model=None, data_dir=args.data_dir
    )
    validation.load_model(model_dir=args.model_dir)
    _, _, _shape = validation.load_images(args.pair_file)
    tpr, fpr, acc, std = validation.verify()
    logging.info(
        "Testing Accuracy {} for {} in shape {}".format(acc, args.pair_file, _shape[0])
    )


[((13, 0, 15, 1), 'logging.basicConfig', 'logging.basicConfig', (), '', False, 'import logging\n'), ((266, 13, 266, 77), 'argparse.ArgumentParser', 'argparse.ArgumentParser', (), '', False, 'import argparse\n'), ((21, 8, 21, 50), 'logging.info', 'logging.info', ({(21, 21, 21, 49): '"""Face Verification for RFW."""'}, {}), "('Face Verification for RFW.')", False, 'import logging\n'), ((28, 8, 28, 37), 'logging.info', 'logging.info', ({(28, 21, 28, 36): '"""Model Loading"""'}, {}), "('Model Loading')", False, 'import logging\n'), ((29, 14, 29, 23), 'mxnet.gpu', 'mx.gpu', ({(29, 21, 29, 22): '0'}, {}), '(0)', True, 'import mxnet as mx\n'), ((30, 38, 30, 76), 'mxnet.model.load_checkpoint', 'mx.model.load_checkpoint', ({(30, 63, 30, 72): 'model_dir', (30, 74, 30, 75): '1'}, {}), '(model_dir, 1)', True, 'import mxnet as mx\n'), ((33, 21, 33, 77), 'mxnet.mod.Module', 'mx.mod.Module', (), '', True, 'import mxnet as mx\n'), ((43, 8, 43, 42), 'logging.info', 'logging.info', ({(43, 21, 43, 41): '"""Image Data Loading"""'}, {}), "('Image Data Loading')", False, 'import logging\n'), ((45, 16, 45, 41), 'pandas.read_csv', 'pd.read_csv', ({(45, 28, 45, 40): 'inp_csv_file'}, {}), '(inp_csv_file)', True, 'import pandas as pd\n'), ((98, 22, 98, 45), 'numpy.asarray', 'np.asarray', ({(98, 33, 98, 44): 'issame_list'}, {}), '(issame_list)', True, 'import numpy as np\n'), ((113, 17, 113, 44), 'mxnet.ndarray.ones', 'nd.ones', ({(113, 25, 113, 43): '(self.batch_size,)'}, {}), '((self.batch_size,))', True, 'from mxnet import ndarray as nd\n'), ((151, 21, 151, 64), 'sklearn.preprocessing.normalize', 'sklearn.preprocessing.normalize', ({(151, 53, 151, 63): 'embeddings'}, {}), '(embeddings)', False, 'import sklearn\n'), ((165, 21, 165, 44), 'numpy.arange', 'np.arange', ({(165, 31, 165, 33): '-1', (165, 35, 165, 36): '1', (165, 38, 165, 43): '0.001'}, {}), '(-1, 1, 0.001)', True, 'import numpy as np\n'), ((186, 17, 186, 68), 'sklearn.model_selection.StratifiedKFold', 'StratifiedKFold', (), '', False, 'from sklearn.model_selection import StratifiedKFold\n'), ((188, 15, 188, 54), 'numpy.zeros', 'np.zeros', ({(188, 24, 188, 53): '(nrof_folds, nrof_thresholds)'}, {}), '((nrof_folds, nrof_thresholds))', True, 'import numpy as np\n'), ((189, 15, 189, 54), 'numpy.zeros', 'np.zeros', ({(189, 24, 189, 53): '(nrof_folds, nrof_thresholds)'}, {}), '((nrof_folds, nrof_thresholds))', True, 'import numpy as np\n'), ((191, 15, 191, 54), 'numpy.zeros', 'np.zeros', ({(191, 24, 191, 53): '(nrof_folds, nrof_thresholds)'}, {}), '((nrof_folds, nrof_thresholds))', True, 'import numpy as np\n'), ((192, 15, 192, 54), 'numpy.zeros', 'np.zeros', ({(192, 24, 192, 53): '(nrof_folds, nrof_thresholds)'}, {}), '((nrof_folds, nrof_thresholds))', True, 'import numpy as np\n'), ((194, 14, 194, 36), 'numpy.zeros', 'np.zeros', ({(194, 24, 194, 34): 'nrof_folds'}, {}), '(nrof_folds)', True, 'import numpy as np\n'), ((196, 19, 196, 41), 'numpy.zeros', 'np.zeros', ({(196, 29, 196, 39): 'nrof_folds'}, {}), '(nrof_folds)', True, 'import numpy as np\n'), ((197, 18, 197, 39), 'numpy.arange', 'np.arange', ({(197, 28, 197, 38): 'nrof_pairs'}, {}), '(nrof_pairs)', True, 'import numpy as np\n'), ((199, 18, 199, 68), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((200, 16, 200, 40), 'numpy.mean', 'np.mean', (), '', True, 'import numpy as np\n'), ((203, 15, 203, 56), 'numpy.sum', 'np.sum', (), '', True, 'import numpy as np\n'), ((246, 25, 246, 49), 'numpy.less', 'np.less', ({(246, 33, 246, 37): 'dist', (246, 39, 246, 48): 'threshold'}, {}), '(dist, threshold)', True, 'import numpy as np\n'), ((247, 24, 247, 51), 'numpy.less', 'np.less', ({(247, 32, 247, 45): 'actual_issame', (247, 47, 247, 50): '0.5'}, {}), '(actual_issame, 0.5)', True, 'import numpy as np\n'), ((260, 13, 260, 68), 'sklearn.metrics.f1_score', 'sklearn.metrics.f1_score', ({(260, 38, 260, 52): 'predict_issame', (260, 54, 260, 67): 'actual_issame'}, {}), '(predict_issame, actual_issame)', False, 'import sklearn\n'), ((50, 19, 50, 86), 'mxnet.ndarray.empty', 'nd.empty', ({(50, 28, 50, 85): '(pairs.shape[0] * 2, 3, self.image_size, self.image_size)'}, {}), '((pairs.shape[0] * 2, 3, self.image_size, self.image_size))', True, 'from mxnet import ndarray as nd\n'), ((72, 18, 72, 35), 'cv2.imread', 'cv2.imread', ({(72, 29, 72, 34): 'path1'}, {}), '(path1)', False, 'import cv2\n'), ((73, 18, 73, 54), 'cv2.cvtColor', 'cv2.cvtColor', ({(73, 31, 73, 34): 'im1', (73, 36, 73, 53): 'cv2.COLOR_BGR2RGB'}, {}), '(im1, cv2.COLOR_BGR2RGB)', False, 'import cv2\n'), ((74, 18, 74, 46), 'numpy.transpose', 'np.transpose', ({(74, 31, 74, 34): 'im1', (74, 36, 74, 45): '(2, 0, 1)'}, {}), '(im1, (2, 0, 1))', True, 'import numpy as np\n'), ((75, 18, 75, 34), 'mxnet.nd.array', 'mx.nd.array', ({(75, 30, 75, 33): 'im1'}, {}), '(im1)', True, 'import mxnet as mx\n'), ((77, 18, 77, 35), 'cv2.imread', 'cv2.imread', ({(77, 29, 77, 34): 'path2'}, {}), '(path2)', False, 'import cv2\n'), ((78, 18, 78, 54), 'cv2.cvtColor', 'cv2.cvtColor', ({(78, 31, 78, 34): 'im2', (78, 36, 78, 53): 'cv2.COLOR_BGR2RGB'}, {}), '(im2, cv2.COLOR_BGR2RGB)', False, 'import cv2\n'), ((79, 18, 79, 46), 'numpy.transpose', 'np.transpose', ({(79, 31, 79, 34): 'im2', (79, 36, 79, 45): '(2, 0, 1)'}, {}), '(im2, (2, 0, 1))', True, 'import numpy as np\n'), ((80, 18, 80, 34), 'mxnet.nd.array', 'mx.nd.array', ({(80, 30, 80, 33): 'im2'}, {}), '(im2)', True, 'import mxnet as mx\n'), ((159, 21, 159, 38), 'numpy.mean', 'np.mean', ({(159, 29, 159, 37): 'accuracy'}, {}), '(accuracy)', True, 'import numpy as np\n'), ((159, 40, 159, 56), 'numpy.std', 'np.std', ({(159, 47, 159, 55): 'accuracy'}, {}), '(accuracy)', True, 'import numpy as np\n'), ((172, 12, 172, 37), 'numpy.asarray', 'np.asarray', ({(172, 23, 172, 36): 'actual_issame'}, {}), '(actual_issame)', True, 'import numpy as np\n'), ((204, 55, 204, 85), 'numpy.linalg.norm', 'line.norm', (), '', True, 'from numpy import linalg as line\n'), ((213, 24, 213, 51), 'numpy.zeros', 'np.zeros', ({(213, 34, 213, 49): 'nrof_thresholds'}, {}), '(nrof_thresholds)', True, 'import numpy as np\n'), ((218, 35, 218, 55), 'numpy.argmax', 'np.argmax', ({(218, 45, 218, 54): 'acc_train'}, {}), '(acc_train)', True, 'import numpy as np\n'), ((237, 14, 237, 30), 'numpy.mean', 'np.mean', ({(237, 22, 237, 26): 'tprs', (237, 28, 237, 29): '(0)'}, {}), '(tprs, 0)', True, 'import numpy as np\n'), ((238, 14, 238, 30), 'numpy.mean', 'np.mean', ({(238, 22, 238, 26): 'fprs', (238, 28, 238, 29): '(0)'}, {}), '(fprs, 0)', True, 'import numpy as np\n'), ((121, 24, 121, 87), 'mxnet.ndarray.slice_axis', 'nd.slice_axis', (), '', True, 'from mxnet import ndarray as nd\n'), ((122, 24, 122, 47), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((124, 21, 124, 68), 'mxnet.io.DataBatch', 'mx.io.DataBatch', (), '', True, 'import mxnet as mx\n'), ((129, 27, 129, 50), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((143, 24, 143, 43), 'numpy.linalg.norm', 'np.linalg.norm', ({(143, 39, 143, 42): '_em'}, {}), '(_em)', True, 'import numpy as np\n'), ((204, 22, 204, 52), 'numpy.linalg.norm', 'line.norm', (), '', True, 'from numpy import linalg as line\n'), ((249, 25, 251, 9), 'sklearn.metrics.confusion_matrix', 'sklearn.metrics.confusion_matrix', ({(250, 12, 250, 25): 'actual_issame', (250, 27, 250, 41): 'predict_issame'}, {}), '(actual_issame, predict_issame)', False, 'import sklearn\n'), ((85, 26, 85, 49), 'mxnet.ndarray.flip', 'mx.ndarray.flip', ({(85, 42, 85, 45): 'im1', (85, 47, 85, 48): '2'}, {}), '(im1, 2)', True, 'import mxnet as mx\n'), ((90, 26, 90, 49), 'mxnet.ndarray.flip', 'mx.ndarray.flip', ({(90, 42, 90, 45): 'im2', (90, 47, 90, 48): '2'}, {}), '(im2, 2)', True, 'import mxnet as mx\n'), ((133, 33, 133, 80), 'numpy.zeros', 'np.zeros', ({(133, 42, 133, 79): '(data.shape[0], _embeddings.shape[1])'}, {}), '((data.shape[0], _embeddings.shape[1]))', True, 'import numpy as np\n')]