import tempfile
import unittest
from pathlib import Path

import laspy
import numpy as np

from B04_wf1_Surface.B04_wf1_Surface_Engine_Structurize import structurize_las


class StructurizeLasTest(unittest.TestCase):
    def test_structurize_las(self) -> None:
        with tempfile.TemporaryDirectory() as temporary_directory:
            source = Path(temporary_directory) / "source.las"
            output_dir = Path(temporary_directory) / "processed"
            las = laspy.LasData(laspy.LasHeader(point_format=3, version="1.2"))
            las.x = np.array([100.0, 102.0])
            las.y = np.array([200.0, 204.0])
            las.z = np.array([10.0, 14.0])
            las.intensity = np.array([7, 9], dtype=np.uint16)
            las.red = np.array([65535, 256], dtype=np.uint16)
            las.green = np.array([32768, 512], dtype=np.uint16)
            las.blue = np.array([0, 768], dtype=np.uint16)
            las.classification = np.array([2, 5], dtype=np.uint8)
            las.write(source)
            progress: list[int] = []

            target = structurize_las(source, output_dir, progress.append)

            self.assertEqual(target, output_dir / "structured.npz")
            self.assertEqual(progress[-1], 100)
            with np.load(target) as structured:
                np.testing.assert_allclose(
                    structured["xyz"],
                    np.array([[100.0, 200.0, 10.0], [102.0, 204.0, 14.0]]),
                )
                np.testing.assert_array_equal(structured["intensity"], [7, 9])
                np.testing.assert_array_equal(structured["rgb"], [[255, 128, 0], [1, 2, 3]])
                np.testing.assert_array_equal(structured["classification"], [2, 5])
                np.testing.assert_array_equal(structured["total_points"], [2])