from __future__ import annotations

import copy
import math

import numpy as np
import open3d as o3d

from preprocessing.plane import PlaneResult


class Visualizer:
    """Open3D visualization utilities."""

    def show_registration(
        self,
        source: o3d.geometry.PointCloud,
        target: o3d.geometry.PointCloud,
        recovered: o3d.geometry.PointCloud,
    ) -> None:
        source_colored = copy.deepcopy(source)
        target_colored = copy.deepcopy(target)
        recovered_colored = copy.deepcopy(recovered)
        source_colored.paint_uniform_color([0, 0, 1])
        target_colored.paint_uniform_color([0, 1, 0])
        recovered_colored.paint_uniform_color([1, 0, 0])

        o3d.visualization.draw_geometries(
            [source_colored, target_colored],
            window_name="Before Registration (Blue: Source, Green: Target)",
            width=1024,
            height=768,
        )
        o3d.visualization.draw_geometries(
            [recovered_colored, target_colored],
            window_name="After Registration (Red: Recovered, Green: Target)",
            width=1024,
            height=768,
        )

    def show_denoise(
        self,
        clean: o3d.geometry.PointCloud,
        noisy: o3d.geometry.PointCloud,
        filtered: o3d.geometry.PointCloud | None = None,
        filter_name: str = "Filtered",
    ) -> None:
        clean_colored = copy.deepcopy(clean)
        noisy_colored = copy.deepcopy(noisy)
        clean_colored.paint_uniform_color([0, 0, 1])
        noisy_colored.paint_uniform_color([1.0, 0.7, 0.5])
        o3d.visualization.draw_geometries(
            [clean_colored, noisy_colored],
            window_name=f"Clean (Blue, {len(clean.points)}) vs Noisy (Orange, {len(noisy.points)})",
        )

        if filtered is not None:
            filtered_colored = copy.deepcopy(filtered)
            filtered_colored.paint_uniform_color([0, 1, 0])
            o3d.visualization.draw_geometries(
                [clean_colored, filtered_colored],
                window_name=f"{filter_name}: Clean (Blue) vs Filtered (Green)",
            )

    def show_curvature(self, pointcloud: o3d.geometry.PointCloud, window_name: str = "Curvature Visualization") -> None:
        o3d.visualization.draw_geometries([pointcloud], window_name=window_name)

    def show_combined_pointclouds(
        self,
        pointclouds: list[o3d.geometry.PointCloud],
        labels: list[str] | None = None,
        window_name: str = "Combined Visualization",
        arrange_grid: bool = False,
    ) -> None:
        combined = self.combine_pointclouds(pointclouds, arrange_grid=arrange_grid)
        if len(combined.points) == 0:
            return

        title = window_name
        if labels:
            title = f"{window_name} ({len(labels)} results)"
        o3d.visualization.draw_geometries([combined], window_name=title, width=1280, height=900)

    def combine_pointclouds(
        self,
        pointclouds: list[o3d.geometry.PointCloud],
        arrange_grid: bool = False,
    ) -> o3d.geometry.PointCloud:
        geometries = [
            copy.deepcopy(pointcloud)
            for pointcloud in pointclouds
            if pointcloud is not None and len(pointcloud.points) > 0
        ]
        if not geometries:
            return o3d.geometry.PointCloud()

        for geometry in geometries:
            if not geometry.has_colors():
                geometry.paint_uniform_color([0.62, 0.62, 0.62])

        if arrange_grid and len(geometries) > 1:
            extents = [
                np.asarray(geometry.get_max_bound()) - np.asarray(geometry.get_min_bound())
                for geometry in geometries
            ]
            max_extent = max(float(np.max(extent)) for extent in extents)
            spacing = max(max_extent * 1.35, 1.0)
            cols = math.ceil(math.sqrt(len(geometries)))

            for index, geometry in enumerate(geometries):
                row, col = divmod(index, cols)
                center = (np.asarray(geometry.get_min_bound()) + np.asarray(geometry.get_max_bound())) / 2.0
                target_center = np.array([col * spacing, -row * spacing, 0.0])
                geometry.translate(target_center - center)

        combined = o3d.geometry.PointCloud()
        for geometry in geometries:
            combined += geometry
        return combined

    def show_plane(self, plane_result: PlaneResult) -> None:
        pointcloud = copy.deepcopy(plane_result.pointcloud)
        pointcloud.paint_uniform_color([0.2, 0.4, 1.0])
        inlier_cloud = pointcloud.select_by_index(plane_result.inliers)
        inlier_cloud.paint_uniform_color([1.0, 0.2, 0.2])

        plane_mesh = self._create_plane_mesh(plane_result.model)
        o3d.visualization.draw_geometries(
            [pointcloud, inlier_cloud, plane_mesh],
            window_name="Plane Fitting",
            width=1024,
            height=768,
        )

    def _create_plane_mesh(self, plane_model: tuple[float, float, float, float]) -> o3d.geometry.TriangleMesh:
        a, b, c, d = plane_model
        size = 1.0
        xx, yy = np.meshgrid(np.linspace(-size, size, 2), np.linspace(-size, size, 2))
        if abs(c) < 1e-12:
            zz = np.zeros_like(xx)
        else:
            zz = -(a * xx + b * yy + d) / c

        vertices = np.column_stack([xx.ravel(), yy.ravel(), zz.ravel()])
        triangles = np.array([[0, 1, 2], [1, 3, 2]])
        mesh = o3d.geometry.TriangleMesh()
        mesh.vertices = o3d.utility.Vector3dVector(vertices)
        mesh.triangles = o3d.utility.Vector3iVector(triangles)
        mesh.paint_uniform_color([0.1, 0.9, 0.1])
        mesh.compute_vertex_normals()
        return mesh