import os
import pyrealsense2 as rs
import numpy as np

from unittest import TestCase

from ..session import Liveness, CameraParams
from ..utils import estimate_depth_bounding_box, expand_bbox_to_edge_and_crop
from ..types import Rectangle


ASSETS_PATH = os.path.join(os.path.dirname(__file__), "assets")
liveness_session = None
expected_projected_bounding_box = [
    514.62513733,
    267.85284424,
    726.44473267,
    551.44799805,
]


class TestUtils(TestCase):
    @classmethod
    def setUpClass(cls):
        global liveness_session
        liveness_session = Liveness()

    def setUp(self):
        self.liveness_session = liveness_session

    def test_estimate_depth_bounding_box(self):
        bounding_box = Rectangle(
            528.551139831543, 234.36917863815668, 839.0621948242188, 642.2044240250417
        )

        depth_intrinsic = rs.intrinsics()
        depth_intrinsic.width = 1280
        depth_intrinsic.height = 720
        depth_intrinsic.ppx = 640.387
        depth_intrinsic.ppy = 357.513
        depth_intrinsic.fx = 635.811
        depth_intrinsic.fy = 635.811
        depth_intrinsic.model = rs.distortion.brown_conrady
        depth_intrinsic.coeffs = [0, 0, 0, 0, 0]

        color_intrinsic = rs.intrinsics()
        color_intrinsic.width = 1280
        color_intrinsic.height = 720
        color_intrinsic.ppx = 647.024
        color_intrinsic.ppy = 356.927
        color_intrinsic.fx = 922.169
        color_intrinsic.fy = 922.476
        color_intrinsic.model = rs.distortion.inverse_brown_conrady
        color_intrinsic.coeffs = [0, 0, 0, 0, 0]

        color_depth_extr = rs.extrinsics()
        color_depth_extr.rotation = [
            0.999945,
            0.0103263,
            0.00163071,
            -0.0103348,
            0.999932,
            0.00530964,
            -0.00157577,
            -0.0053262,
            0.999985,
        ]
        color_depth_extr.translation = [-0.0147758, -0.000159923, -0.000372309]
        camera_params = CameraParams(depth_intrinsic, color_intrinsic, color_depth_extr)

        proj_depth_bb = estimate_depth_bounding_box(bounding_box, camera_params)

        self.assertIsNotNone(
            proj_depth_bb, msg="unexpected or none projected depth bounding box"
        )

        self.assertTrue(
            np.allclose(proj_depth_bb, expected_projected_bounding_box),
            msg="invalid projected depth bounding box",
        )

    def test_expand_bbox_to_edge_and_crop(self):
        depth_frame = liveness_session.load_depth_data_from_file(
            os.path.join(ASSETS_PATH, "depth.txt")
        )
        expanded_cropped_frame = expand_bbox_to_edge_and_crop(
            depth_frame, expected_projected_bounding_box
        )

        expected_cropped_depth_frame = liveness_session.load_depth_data_from_file(
            os.path.join(ASSETS_PATH, "cropped_depth.txt")
        )

        self.assertIsNotNone(
            expanded_cropped_frame, msg="unexpected cropped depth frame"
        )
        self.assertTrue(
            np.array_equal(expanded_cropped_frame, expected_cropped_depth_frame),
            msg="invalid cropped depth frame",
        )