from __future__ import annotations

import tempfile
import unittest
from pathlib import Path

import numpy as np
import open3d as o3d

import sys

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from run_pointcloud_understanding import classify_semantic_points, process_point_cloud  # noqa: E402


def write_cloud(path: Path, points: np.ndarray) -> None:
    cloud = o3d.geometry.PointCloud()
    cloud.points = o3d.utility.Vector3dVector(points)
    assert o3d.io.write_point_cloud(str(path), cloud)


class PointCloudUnderstandingTests(unittest.TestCase):
    def test_raster_vector_and_mesh_are_created(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            rng = np.random.default_rng(2)
            ground = np.column_stack((rng.uniform(-4, 4, (600, 2)), rng.normal(0, 0.01, 600)))
            roof = np.column_stack((rng.uniform(-1, 1, (350, 2)), rng.normal(2, 0.02, 350)))
            source = root / "site.ply"
            write_cloud(source, np.vstack((ground, roof)))
            result = process_point_cloud(source, root / "output", voxel_size=0.2, ground_distance=0.1, elevated_threshold=0.6)
            self.assertGreater(result["elevated_points"], 0)
            self.assertGreater(result["elevated_footprint_count"], 0)
            self.assertEqual(result["coordinate_basis"], "local_point_cloud_coordinates")
            for field in ("dsm_file", "label_raster_file", "preview_file", "vector_file", "summary_file", "classified_point_cloud"):
                self.assertTrue((root / "output" / result[field]).is_file())

    def test_small_cloud_fails(self) -> None:
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            source = root / "small.ply"
            write_cloud(source, np.zeros((4, 3)))
            with self.assertRaises(ValueError):
                process_point_cloud(source, root / "output")

    def test_semantic_rules_keep_unknown_and_identify_strong_candidates(self) -> None:
        ground = np.column_stack((np.linspace(0, 4, 20), np.zeros(20), np.zeros(20)))
        structure = np.column_stack((10 + np.linspace(0, 0.8, 20), np.zeros(20), np.full(20, 3.0)))
        vegetation = np.column_stack((20 + np.linspace(0, 0.8, 20), np.zeros(20), np.linspace(2.0, 3.5, 20)))
        wire = np.column_stack((30 + np.arange(20), np.zeros(20), np.linspace(6.0, 6.1, 20)))
        pole = np.column_stack((40 + np.zeros(20), np.zeros(20), np.linspace(1.0, 10.0, 20)))
        points = np.vstack((ground, structure, vegetation, wire, pole))
        heights = points[:, 2].copy()
        colors = np.full((len(points), 3), 0.5)
        colors[len(ground) + len(structure):len(ground) + len(structure) + len(vegetation)] = [0.15, 0.7, 0.15]
        labels = classify_semantic_points(points, colors, heights, voxel_size=0.2)
        self.assertIn(2, labels)
        self.assertIn(5, labels)
        self.assertIn(6, labels)
        self.assertIn(15, labels)
        self.assertIn(16, labels)


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