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

import argparse
import numpy as np
import cv2

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


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


def load_image(path, input_details):
    img = cv2.imread(path, cv2.IMREAD_COLOR)
    if img is None:
        raise FileNotFoundError(f"Could not load image: {path}")
    h, w = input_details[0]['shape'][1], input_details[0]['shape'][2]
    img = cv2.resize(img, (w, h))

    # Handle uint8 quantized input
    if input_details[0]['dtype'] == np.uint8:
        img = img.astype(np.uint8)
    else:
        img = (img / 255.0).astype(np.float32)

    return img[None, ...]


def get_embedding(interpreter, input_details, output_details, img):
    interpreter.set_tensor(input_details[0]['index'], img)
    interpreter.invoke()
    out = interpreter.get_tensor(output_details[0]['index'])

    # Dequantize output if needed
    if output_details[0]['dtype'] == np.uint8:
        scale, zero_point = output_details[0]['quantization']
        out = (out.astype(np.float32) - zero_point) * scale

    return out[0]


def main():
    parser = argparse.ArgumentParser(description="FaceNet512 face similarity example")
    parser.add_argument('-m', '--model', default='original_model/facenet512_uint8_float32.tflite',
                        help='Path to the TFLite model file')
    parser.add_argument('-i', '--input', default='face.jpg',
                        help='Path to the first face image')
    parser.add_argument('-i2', '--input2', default='face2.jpg',
                        help='Path to the second face image')
    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()

    img1 = load_image(args.input, input_details)
    img2 = load_image(args.input2, input_details)

    emb1 = get_embedding(interpreter, input_details, output_details, img1)
    emb2 = get_embedding(interpreter, input_details, output_details, img2)

    dist = cosine_similarity(emb1, emb2)
    print(f"Cosine similarity distance: {dist:.4f}")
    print("Same face" if dist < 0.3 else "Different face")


if __name__ == '__main__':
    main()