#!/usr/bin/env python3
# Copyright 2022-2024,2026 NXP
# SPDX-License-Identifier: MIT

import argparse
import os
import numpy as np
from glob import glob
import tqdm

try:
    import tflite_runtime.interpreter as tflite
except ImportError:
    import tensorflow as tf
    tflite = tf.lite

import cv2

THRESHOLD = 0.6

# Download and extract https://vis-www.cs.umass.edu/lfw/lfw-deepfunneled.tgz
LFW_DIR = "lfw-deepfunneled"

# Download https://vis-www.cs.umass.edu/lfw/pairsDevTest.txt
LFW_PAIRS_FILE = "pairsDevTest.txt"


def cosine_similarity(a, b):
    return 1 - np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))


def load_img(filename):
    img = cv2.imread(filename, 1)
    img = img[45:-45, 45:-45]
    img = np.array(img, dtype=np.uint8)
    return img[None, ...]


def main():
    parser = argparse.ArgumentParser(description="Evaluate FaceNet512 on LFW pairsDevTest")
    parser.add_argument('-m', '--model',
                        default='original_model/facenet512_uint8_float32.tflite',
                        help='Path to the TFLite model file')
    args = parser.parse_args()

    interpreter = tflite.Interpreter(model_path=args.model)
    interpreter.allocate_tensors()
    input_details = interpreter.get_input_details()
    output_details = interpreter.get_output_details()

    with open(LFW_PAIRS_FILE, 'r') as f:
        pairs = f.readlines()[1:]
        pairs = [p.strip().split("\t") for p in pairs]

    image_filenames = set()
    for line in pairs[:500]:
        image_filenames.add(os.sep.join([LFW_DIR, line[0],
                                        f"{line[0]}_{int(line[1]):04d}.jpg"]))
        image_filenames.add(os.sep.join([LFW_DIR, line[0],
                                        f"{line[0]}_{int(line[2]):04d}.jpg"]))
    for line in pairs[500:]:
        image_filenames.add(os.sep.join([LFW_DIR, line[0],
                                        f"{line[0]}_{int(line[1]):04d}.jpg"]))
        image_filenames.add(os.sep.join([LFW_DIR, line[2],
                                        f"{line[2]}_{int(line[3]):04d}.jpg"]))

    feature_vectors = dict()
    for f in tqdm.tqdm(image_filenames, desc="Running inferences"):
        img = load_img(f)
        interpreter.set_tensor(input_details[0]['index'], img)
        interpreter.invoke()
        out = interpreter.get_tensor(output_details[0]['index'])
        # Dequantize if needed
        if output_details[0]['dtype'] == np.uint8:
            scale, zero_point = output_details[0]['quantization']
            out = (out.astype(np.float32) - zero_point) * scale
        feature_vectors[f] = out[0]

    n_correct = 0

    # Evaluate same person pairs
    for line in pairs[:500]:
        f1 = os.sep.join([LFW_DIR, line[0], f"{line[0]}_{int(line[1]):04d}.jpg"])
        f2 = os.sep.join([LFW_DIR, line[0], f"{line[0]}_{int(line[2]):04d}.jpg"])
        result = cosine_similarity(feature_vectors[f1], feature_vectors[f2])
        n_correct += 1 if result < THRESHOLD else 0

    # Evaluate different person pairs
    for line in pairs[500:]:
        f1 = os.sep.join([LFW_DIR, line[0], f"{line[0]}_{int(line[1]):04d}.jpg"])
        f2 = os.sep.join([LFW_DIR, line[2], f"{line[2]}_{int(line[3]):04d}.jpg"])
        result = cosine_similarity(feature_vectors[f1], feature_vectors[f2])
        n_correct += 1 if result > THRESHOLD else 0

    print(f"Quantized model accuracy: {n_correct / 1000.0:.1%}")


if __name__ == '__main__':
    main()