import os
import numpy as np
from unittest import TestCase

from ..sdk import SDK
from ..engine import Engine
from ..utils import load_image
from ..types import ScoringMode, Embedding

from .constants import (
    IMG_ONEFACE,
    IMG_IDENTITY1_FACE1,
    IMG_IDENTITY1_FACE2,
    IMG_ONEFACE_RECO_INPUT_IMG,
    ERR_UNEXPECTED_NUM_FACES,
    ERR_UNEXPECTED_NUM_INFERENCES,
    ERR_MISSING_EMBEDDING,
    ERR_INVALID_SCORING_MODE,
    ERR_INVALID_EMBEDDING_SIZE,
    EXPECTED_STANDARD_EMBED_LEN,
)

ASSETS_PATH = os.path.join(os.path.dirname(__file__), "assets")
engine_default = None
scoring_mode = None
sdk = None


class TestSDK(TestCase):
    @classmethod
    def setUpClass(cls):
        global sdk
        global engine_default
        global scoring_mode

        engine_default = Engine.OPENVINO
        scoring_mode = ScoringMode.StandardEmbedding

        sdk = SDK(engine=engine_default, settings={"scoring_mode": scoring_mode})

    def setUp(self):
        self.sdk = sdk

    def test_get_faces_with_embeddings(self):
        imgs = [load_image(os.path.join(ASSETS_PATH, IMG_ONEFACE))]
        detection_result = self.sdk.get_faces(imgs, embeddings=True)

        faces = detection_result.faces
        self.assertEqual(len(faces), 1, msg=ERR_UNEXPECTED_NUM_FACES)

        image_inferences = detection_result.image_inferences
        self.assertEqual(len(image_inferences), 1, msg=ERR_UNEXPECTED_NUM_INFERENCES)

        self.assert_embeddings(faces)

    def test_get_embedding_from_landmarks(self):
        imgs = [load_image(os.path.join(ASSETS_PATH, IMG_ONEFACE))]
        detection_result = self.sdk.get_faces(imgs, embeddings=True)

        faces = detection_result.faces
        f = faces[0]
        landmarks = f.landmarks

        embeddings = self.sdk.get_embeddings_from_landmarks(imgs[0], [landmarks, landmarks])
        self.assertEqual(len(embeddings), 2)
        embedding = embeddings[0]

        self.assertTrue(embedding.scoring_mode == ScoringMode.StandardEmbedding)
        similarity = SDK.get_similarity(f.embedding, embedding)
        self.assertAlmostEqual(similarity, 4.0, delta=0.01)

    def test_scoring_same_image(self):
        img = load_image(os.path.join(ASSETS_PATH, IMG_ONEFACE))
        faces = self.sdk.get_faces([img, img], embeddings=True).faces

        similarity = SDK.get_similarity(faces[0].embedding, faces[1].embedding)
        self.assertAlmostEqual(similarity, 4.0, delta=0.01)

        confidence = self.sdk.get_confidence(faces[0].embedding, faces[1].embedding)
        self.assertAlmostEqual(confidence, 1.0, delta=0.01)

        match_score = SDK.get_match_score(faces[0].embedding, faces[1].embedding)
        self.assertAlmostEqual(match_score, 957, delta=1)

    def test_scoring_same_identity(self):
        img1 = load_image(os.path.join(ASSETS_PATH, IMG_IDENTITY1_FACE1))
        img2 = load_image(os.path.join(ASSETS_PATH, IMG_IDENTITY1_FACE2))
        faces = self.sdk.get_faces([img1, img2], embeddings=True).faces

        similarity = SDK.get_similarity(faces[0].embedding, faces[1].embedding)
        self.assertAlmostEqual(similarity, 3.58, delta=0.01)

        confidence = self.sdk.get_confidence(faces[0].embedding, faces[1].embedding)
        self.assertAlmostEqual(confidence, 1.0, delta=0.01)

        match_score = SDK.get_match_score(faces[0].embedding, faces[1].embedding)
        self.assertAlmostEqual(match_score, 903, delta=2)

    def test_scoring_diff_identity(self):
        img1 = load_image(os.path.join(ASSETS_PATH, IMG_IDENTITY1_FACE1))
        img2 = load_image(os.path.join(ASSETS_PATH, IMG_ONEFACE))
        faces = self.sdk.get_faces([img1, img2], embeddings=True).faces

        similarity = SDK.get_similarity(faces[0].embedding, faces[1].embedding)
        self.assertAlmostEqual(similarity, 1.85, delta=0.01)

        confidence = self.sdk.get_confidence(faces[0].embedding, faces[1].embedding)
        self.assertAlmostEqual(confidence, 0, delta=0.01)

        match_score = SDK.get_match_score(faces[0].embedding, faces[1].embedding)
        self.assertAlmostEqual(match_score, 198, delta=2)

    def test_get_embedding_from_prepared_image(self):
        imgs = [load_image(os.path.join(ASSETS_PATH, IMG_ONEFACE))]
        detection_result = self.sdk.get_faces(
            imgs, qualities=True, landmarks=True, embeddings=True
        )
        faces = detection_result.faces
        f = faces[0]

        reco_img = load_image(os.path.join(ASSETS_PATH, IMG_ONEFACE_RECO_INPUT_IMG))
        embedding = self.sdk.get_embedding_from_prepared_image(reco_img)
        self.assertTrue(len(embedding.data) == EXPECTED_STANDARD_EMBED_LEN)
        self.assertTrue(embedding.scoring_mode == scoring_mode)
        self.assertTrue(
            np.allclose(f.embedding.data, embedding.data, rtol=0, atol=0.001),
            msg="Invalid embedding value",
        )

    def assert_embeddings(self, faces):
        for f in faces:
            self.assertIsNotNone(f.embedding, msg=ERR_MISSING_EMBEDDING)
            self.assertEqual(
                f.embedding.scoring_mode,
                ScoringMode.StandardEmbedding,
                msg=ERR_INVALID_SCORING_MODE,
            )
            self.assertTrue(
                len(f.embedding.data) in Embedding.STANDARD_SIZES,
                msg=ERR_INVALID_EMBEDDING_SIZE,
            )