from __future__ import annotations

import copy

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_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