import unittest
import numpy as np

import importlib
utils = importlib.import_module('extensions.sd-webui-controlnet.tests.utils', 'utils')
utils.setup_test_env()

from annotator.openpose.util import faceDetect, handDetect
from annotator.openpose.body import Keypoint, BodyResult

class TestFaceDetect(unittest.TestCase):
    def test_no_faces(self):
        oriImg = np.zeros((100, 100, 3), dtype=np.uint8)
        body = BodyResult([None] * 18, total_score=3, total_parts=0)
        expected_result = None
        result = faceDetect(body, oriImg)

        self.assertEqual(result, expected_result)

    def test_single_face(self):
        body = BodyResult([
            Keypoint(50, 50),
            *([None] * 13),
            Keypoint(30, 40),
            Keypoint(70, 40),
            Keypoint(20, 50),
            Keypoint(80, 50),
        ], total_score=2, total_parts=5)

        oriImg = np.zeros((100, 100, 3), dtype=np.uint8)

        expected_result = (0, 0, 120)
        result = faceDetect(body, oriImg)

        self.assertEqual(result, expected_result)

class TestHandDetect(unittest.TestCase):
    def test_no_hands(self):
        oriImg = np.zeros((100, 100, 3), dtype=np.uint8)
        body = BodyResult([None] * 18, total_score=3, total_parts=0)
        expected_result = []
        result = handDetect(body, oriImg)

        self.assertEqual(result, expected_result)

    def test_single_left_hand(self):
        oriImg = np.zeros((100, 100, 3), dtype=np.uint8)

        body = BodyResult([
            None, None, None, None, None,
            Keypoint(20, 20),
            Keypoint(40, 30),
            Keypoint(60, 40),
            *([None] * 8),
            Keypoint(20, 60),
            Keypoint(40, 70),
            Keypoint(60, 80)
        ], total_score=3, total_parts=0.5)

        expected_result = [(49, 26, 33, True)]
        result = handDetect(body, oriImg)

        self.assertEqual(result, expected_result)

    def test_single_right_hand(self):
        oriImg = np.zeros((100, 100, 3), dtype=np.uint8)

        body = BodyResult([
            None, None,
            Keypoint(20, 20),
            Keypoint(40, 30),
            Keypoint(60, 40),
            *([None] * 11),
            Keypoint(20, 60),
            Keypoint(40, 70),
            Keypoint(60, 80)
        ], total_score=3, total_parts=0.5)

        expected_result = [(49, 26, 33, False)]
        result = handDetect(body, oriImg)

        self.assertEqual(result, expected_result)

    def test_multiple_hands(self):
        body = BodyResult([
            Keypoint(20, 20),
            Keypoint(40, 30),
            Keypoint(60, 40),
            Keypoint(20, 60),
            Keypoint(40, 70),
            Keypoint(60, 80),
            Keypoint(10, 10),
            Keypoint(30, 20),
            Keypoint(50, 30),
            Keypoint(10, 50),
            Keypoint(30, 60),
            Keypoint(50, 70),
            *([None] * 6),
        ], total_score=3, total_parts=0.5)

        oriImg = np.zeros((100, 100, 3), dtype=np.uint8)

        expected_result = [(0, 0, 100, True), (16, 43, 56, False)]
        result = handDetect(body, oriImg)
        self.assertEqual(result, expected_result)


if __name__ == '__main__':
    unittest.main()