from __future__ import annotations

import sys
import json
import tempfile
import unittest
from pathlib import Path

import numpy as np
import open3d as o3d

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from prepare_multiview_point_features import FEATURE_NAMES, add_annotation_source_from_existing_dataset, photo_features, z_buffer_visible_indices  # noqa: E402


class MultiviewPointFeatureTests(unittest.TestCase):
    def test_z_buffer_keeps_nearest_point_per_two_pixel_cell(self) -> None:
        pixels = np.array([[2.0, 2.0], [3.0, 3.0], [8.0, 4.0]])
        depths = np.array([8.0, 2.0, 3.0])
        visible = z_buffer_visible_indices(pixels, depths, np.array([True, True, True]), width=10, cell_size=2)
        np.testing.assert_array_equal(visible, [1, 2])

    def test_photo_features_are_bounded_and_have_stable_contract(self) -> None:
        photo = np.array([[[0, 0, 0], [255, 0, 0]], [[0, 255, 0], [0, 0, 255]]], dtype=np.uint8)
        values = photo_features(photo)
        self.assertEqual(values.shape, (2, 2, len(FEATURE_NAMES)))
        self.assertTrue(np.isfinite(values).all())
        self.assertTrue(((values >= 0.0) & (values <= 1.0)).all())

    def test_annotation_source_preserves_dataset_order_and_original_rgb(self) -> None:
        with tempfile.TemporaryDirectory() as temp_dir:
            output = Path(temp_dir)
            xyz = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float64)
            rgb = np.array([[0.1, 0.2, 0.3], [0.8, 0.7, 0.6]], dtype=np.float32)
            np.savez_compressed(output / "multiview-point-features.npz", xyz=xyz, las_rgb=rgb)
            (output / "run_metadata.json").write_text(json.dumps({"capability": "05-3d-pointcloud", "artifacts": {}}), encoding="utf-8")

            result = add_annotation_source_from_existing_dataset(output)
            source = output / "multiview-annotation-source.ply"
            cloud = o3d.io.read_point_cloud(str(source))
            metadata = json.loads((output / "run_metadata.json").read_text(encoding="utf-8"))

            self.assertEqual(result["point_count"], 2)
            np.testing.assert_allclose(np.asarray(cloud.points), xyz)
            np.testing.assert_allclose(np.asarray(cloud.colors), rgb, atol=1 / 255)
            self.assertEqual(metadata["annotation_source"]["point_count"], 2)
            self.assertEqual(metadata["annotation_source"]["point_cloud"], source.name)
            self.assertEqual(metadata["annotation_source"]["feature_dataset"], "multiview-point-features.npz")


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