# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Test eye features generator."""

import unittest
import numpy as np
import pytest
from nvidia_tao_tf1.cv.common.dataio.eye_features_generator import EyeFeaturesGenerator


class EyeFeaturesGeneratorTest(unittest.TestCase):
    """Test EyeFeaturesGenerator."""

    # 104 (x, y) coordinates
    # Landmarks from set/germany-1-gaze-1/Data/VdChfpnvDBu7MByIYFUCwM video1_15_750_vc00_07.png
    landmarks_2D = np.asarray([
        240.4823455810547,
        551.151123046875,
        249.34622192382812,
        591.03857421875,
        259.3180847167969,
        628.1560668945312,
        273.7218933105469,
        667.489501953125,
        293.04754638671875,
        704.656982421875,
        318.59527587890625,
        732.306640625,
        346.848876953125,
        752.8043212890625,
        383.3990783691406,
        778.6015625,
        421.562255859375,
        781.989013671875,
        459.86334228515625,
        766.6541748046875,
        487.89031982421875,
        733.9591064453125,
        514.6818237304688,
        702.6900024414062,
        532.0092163085938,
        668.2235717773438,
        543.92333984375,
        634.5968017578125,
        550.9884033203125,
        603.5624389648438,
        553.140869140625,
        579.9955444335938,
        555.6583862304688,
        547.2760009765625,
        272.3310241699219,
        488.6089172363281,
        299.0483703613281,
        477.6076354980469,
        323.4083251953125,
        471.3211975097656,
        356.01922607421875,
        475.64312744140625,
        377.2359619140625,
        489.0018005371094,
        445.6009826660156,
        482.3224792480469,
        472.71124267578125,
        470.535400390625,
        495.1066589355469,
        467.3921813964844,
        514.7517700195312,
        474.8573303222656,
        530.0750122070312,
        489.78759765625,
        413.2937927246094,
        548.6087036132812,
        416.3590087890625,
        571.7369995117188,
        419.1455383300781,
        594.0293579101562,
        422.4565124511719,
        615.2787475585938,
        376.52923583984375,
        629.8828735351562,
        398.4658508300781,
        637.195068359375,
        420.204833984375,
        639.1713256835938,
        441.1532897949219,
        636.404541015625,
        457.7870178222656,
        625.8230590820312,
        315.46246337890625,
        529.6859130859375,
        325.73907470703125,
        524.547607421875,
        370.7980651855469,
        526.52392578125,
        374.3553466796875,
        535.4171142578125,
        366.64788818359375,
        536.010009765625,
        325.34381103515625,
        533.8361206054688,
        454.6097412109375,
        531.4393920898438,
        463.1083984375,
        519.6356811523438,
        496.1587219238281,
        516.6847534179688,
        502.768798828125,
        522.1144409179688,
        497.8112487792969,
        526.7178955078125,
        464.0527038574219,
        529.07861328125,
        370.99566650390625,
        690.7520141601562,
        389.770263671875,
        676.1276245117188,
        403.9993896484375,
        670.7916870117188,
        422.3347473144531,
        672.0076293945312,
        434.2623291015625,
        667.4786376953125,
        448.6631164550781,
        670.0011596679688,
        466.8448181152344,
        681.6611328125,
        454.66357421875,
        697.7672119140625,
        440.85565185546875,
        706.019287109375,
        426.598876953125,
        706.5772094726562,
        410.03192138671875,
        709.9840698242188,
        391.35125732421875,
        707.155029296875,
        389.770263671875,
        687.9852294921875,
        406.3709411621094,
        685.4160766601562,
        423.9597473144531,
        685.6137084960938,
        436.4102478027344,
        683.0445556640625,
        451.2322692871094,
        683.0445556640625,
        437.1334228515625,
        687.640380859375,
        424.5090026855469,
        689.529052734375,
        407.2795715332031,
        689.8966064453125,
        348.108154296875,
        525.959228515625,
        351.7036437988281,
        522.293212890625,
        355.5106506347656,
        525.3952026367188,
        352.1381530761719,
        528.5972290039062,
        481.3041076660156,
        521.7062377929688,
        483.9189758300781,
        518.29296875,
        488.6904296875,
        521.447998046875,
        484.9493408203125,
        524.9349975585938,
        245.04425048828125,
        623.6600341796875,
        227.25071716308594,
        515.8814697265625,
        548.7801513671875,
        612.8729248046875,
        556.8007202148438,
        507.0892639160156,
        350.721923828125,
        515.5374755859375,
        352.60125732421875,
        536.3585815429688,
        351.1640625,
        519.3336181640625,
        356.20318603515625,
        520.873046875,
        359.4877624511719,
        524.7459106445312,
        357.0066223144531,
        530.8558959960938,
        352.21978759765625,
        532.0278930664062,
        347.2783508300781,
        530.6721801757812,
        344.4820861816406,
        526.685302734375,
        347.42608642578125,
        521.79736328125,
        351.86749267578125,
        525.564453125,
        483.0555725097656,
        511.1482849121094,
        486.0395202636719,
        531.6798706054688,
        483.69964599609375,
        514.4667358398438,
        488.69964599609375,
        516.5667114257812,
        492.0996398925781,
        521.3167114257812,
        490.4996643066406,
        525.0667114257812,
        485.5496520996094,
        528.3167114257812,
        480.2496643066406,
        525.8167114257812,
        477.0496520996094,
        522.11669921875,
        479.5996398925781,
        516.8167114257812,
        484.44964599609375,
        521.4166870117188,
        344.7006530761719,
        523.4890747070312,
        476.96893310546875,
        516.9593505859375
    ]).reshape((-1, 2))

    def test_get_eye_features_with_pupils(self):
        # Eye features from Ground_Truth_DataFactory_v2/GT_combined/eye_feature_combined.txt
        expected_eye_features = [
            0.00000000000000000,
            0.00000000000000000,
            0.18428853505018536,
            -0.24571804673358047,
            0.86001316356753166,
            -0.30714755841697561,
            0.98287218693432188,
            -0.18428853505018536,
            0.88048966746199675,
            -0.10238251947232520,
            0.20476503894465040,
            -0.04095300778893008,
            0.61429511683395122,
            -0.20476503894465040,
            30.00000000000000000,
            -10.00000000000000000,
            21.00000000000000000,
            2.00000000000000000,
            -12.00000000000000000,
            5.00000000000000000,
            -18.00000000000000000,
            -1.00000000000000000,
            -13.00000000000000000,
            -5.00000000000000000,
            20.00000000000000000,
            -8.00000000000000000,
            0.58064516129032262,
            1.66666666666666674,
            0.00000000000000000,
            0.00000000000000000,
            0.16862183682891940,
            -0.08431091841445970,
            0.92742010255905660,
            -0.05058655104867581,
            0.99486883729062436,
            0.10117310209735163,
            0.85997136782748884,
            0.11803528578024357,
            0.16862183682891940,
            0.06744873473156776,
            0.61546970442555582,
            -0.06744873473156776,
            36.50000000000000000,
            -4.00000000000000000,
            26.50000000000000000,
            1.00000000000000000,
            -18.50000000000000000,
            -1.00000000000000000,
            -22.50000000000000000,
            -10.00000000000000000,
            -14.50000000000000000,
            -11.00000000000000000,
            26.50000000000000000,
            -8.00000000000000000,
            0.59999999999999998,
            5.50000000000000000]

        pupil_eye_features = EyeFeaturesGenerator(self.landmarks_2D, 104).get_eye_features()

        self.assertEqual(len(pupil_eye_features), len(expected_eye_features))
        for i in range(len(pupil_eye_features)):
            assert pupil_eye_features[i] == pytest.approx(expected_eye_features[i], 1.0)

    def test_get_eye_features_without_pupils(self):
        expected_eye_features = [
            0.00000000000000000,
            0.00000000000000000,
            0.18428853505018536,
            -0.24571804673358047,
            0.86001316356753166,
            -0.30714755841697561,
            0.98287218693432188,
            -0.18428853505018536,
            0.88048966746199675,
            -0.10238251947232520,
            0.20476503894465040,
            -0.04095300778893008,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            0.00000000000000000,
            0.00000000000000000,
            0.16862183682891940,
            -0.08431091841445970,
            0.92742010255905660,
            -0.05058655104867581,
            0.99486883729062436,
            0.10117310209735163,
            0.85997136782748884,
            0.11803528578024357,
            0.16862183682891940,
            0.06744873473156776,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1,
            -1]

        no_pupil_eye_features = EyeFeaturesGenerator(self.landmarks_2D, 68).get_eye_features()

        self.assertEqual(len(no_pupil_eye_features), len(expected_eye_features))
        for i in range(len(no_pupil_eye_features)):
            assert no_pupil_eye_features[i] == pytest.approx(expected_eye_features[i], 1.0)
