from __future__ import annotations

import copy
import numpy as np
import open3d as o3d


class CurvatureEstimator:
    """Estimate local point-cloud curvature using neighborhood covariance."""

    def compute(self, pointcloud: o3d.geometry.PointCloud, radius: float = 2.0, max_nn: int = 30) -> np.ndarray:
        pointcloud.estimate_normals(o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=max_nn))
        points = np.asarray(pointcloud.points)
        kdtree = o3d.geometry.KDTreeFlann(pointcloud)
        curvatures = np.zeros(len(points), dtype=np.float64)

        for i, point in enumerate(points):
            _, indices, _ = kdtree.search_hybrid_vector_3d(point, radius, max_nn)
            if len(indices) < 3:
                curvatures[i] = 1.0
                continue

            neighbors = points[indices]
            cov = np.cov((neighbors - neighbors.mean(axis=0)).T)
            eigenvalues = np.linalg.eigvalsh(cov)
            denom = eigenvalues.sum()
            curvatures[i] = float(eigenvalues[0] / denom) if denom > 0 else 0.0

        return curvatures

    def colorize(self, pointcloud: o3d.geometry.PointCloud, curvatures: np.ndarray, color_scale: float = 3.0) -> o3d.geometry.PointCloud:
        pcd = copy.deepcopy(pointcloud)
        normalized = curvatures / (np.max(curvatures) + 1e-12)
        normalized = np.clip(normalized * color_scale, 0, 1)
        colors = np.zeros((len(normalized), 3))
        colors[:, 0] = normalized
        colors[:, 1] = 1 - normalized
        pcd.colors = o3d.utility.Vector3dVector(colors)
        return pcd