#!/usr/bin/env python

import unittest
import math

from argaze import PupilFeatures
from argaze.PupilAnalysis import WorkloadIndex

class TestWorkloadIndexClass(unittest.TestCase):
    """Test WorkloadIndex class."""

    def test_analysis(self):
        """Test WorkloadIndex analysis."""

        ts_pupil_diameters = {
            0: PupilFeatures.PupilDiameter(1.),
            1: PupilFeatures.PupilDiameter(1.1),
            2: PupilFeatures.PupilDiameter(1.2),
            3: PupilFeatures.PupilDiameter(1.3),
            4: PupilFeatures.PupilDiameter(1.2),
            5: PupilFeatures.PupilDiameter(1.1),
            6: PupilFeatures.PupilDiameter(1.),
            7: PupilFeatures.PupilDiameter(0.9),
            8: PupilFeatures.PupilDiameter(0.8),
            9: PupilFeatures.PupilDiameter(0.7)
        }

        pupil_diameter_analyzer = WorkloadIndex.PupilDiameterAnalyzer(reference=PupilFeatures.PupilDiameter(1.), period=3)
        ts_analysis = pupil_diameter_analyzer.browse(PupilFeatures.TimeStampedPupilDiameters(ts_pupil_diameters))

        # Check result size
        self.assertEqual(len(ts_analysis), 3)

        # Check each workload index
        ts_1, analysis_1 = ts_analysis.pop_first()
        self.assertTrue(math.isclose(analysis_1, 0.1, abs_tol=1e-2))

        ts_2, analysis_2 = ts_analysis.pop_first()
        self.assertTrue(math.isclose(analysis_2, 0.2, abs_tol=1e-2))

        ts_3, analysis_3 = ts_analysis.pop_first()
        self.assertTrue(math.isclose(analysis_3, 0.1, abs_tol=1e-2))
    
if __name__ == '__main__':

    unittest.main()