import numpy as np
from scipy.spatial import cKDTree
import open3d as o3d


class LRFEstimatorPCA:
    """
    Deterministic, index-order–independent Local Reference Frames (LRFs) for 3D point clouds.

    Key design choices:
    - Neighbors ordered by (distance, index) to ensure geometric ordering with deterministic tiebreaks.
    - Normal orientation disambiguated deterministically w.r.t. neighbor-centroid direction.
      * normal_toward_centroid=True  -> inward (toward centroid)
      * normal_toward_centroid=False -> outward (away from centroid)
    - Tangent x-axis from PCA of neighbors projected onto the tangent plane (index-invariant).
    - y = z × x, right-handedness enforced; per-point SVD orthonormalization; global mean-z ≥ 0.
    """

    def __init__(
        self,
        points: np.ndarray,
        indices: np.ndarray = None,
        k: int = 20,
        radius: int = 20,
        use_ball_neighbors_within_radius: bool = False,
        normal_toward_centroid: bool = True,
        enforce_global_z_positive: bool = True,
        degeneracy_eps: float = 1e-8,
    ):
        if not isinstance(points, np.ndarray) or points.ndim != 2 or points.shape[1] != 3:
            raise ValueError("Input `points` must be an (N, 3) NumPy array.")
        self.points = points.astype(np.float64)
        self.N = self.points.shape[0]
        if self.N == 0:
            raise ValueError("Empty point set.")
        self.radius = radius
        # k within [1, N-1]
        self.k = max(1, min(int(k), max(1, self.N - 1)))

        self.use_ball_neighbors_within_radius = bool(use_ball_neighbors_within_radius)
        self.normal_toward_centroid = bool(normal_toward_centroid)
        self.enforce_global_z_positive = bool(enforce_global_z_positive)
        self.degeneracy_eps = float(degeneracy_eps)

        # Optional subset of point indices to compute LRFs for. If None,
        # LRFs are computed for all points.
        if indices is None:
            self.indices = None
        else:
            idxs = np.asarray(indices, dtype=int)
            if idxs.ndim != 1:
                raise ValueError("`indices` must be a 1-D array of integer indices.")
            if np.any(idxs < 0) or np.any(idxs >= self.N):
                raise IndexError("`indices` contains out-of-range values for the point cloud.")
            self.indices = idxs

        self.lrf = None

    def _build_neighbors(self):
        """
        Returns:
            neigh_idx: (N, k) int array of neighbor indices (excluding self),
                       ordered by (distance, then index) for determinism.
        """
        tree = cKDTree(self.points)

        # Build for either all points or the provided subset in `self.indices`.
        if self.indices is None:
            idxs_to_process = np.arange(self.N)
        else:
            idxs_to_process = self.indices

        neigh_idx = [None] * len(idxs_to_process)
        neigh_dist = [None] * len(idxs_to_process)

        for out_i, i in enumerate(idxs_to_process):
            if self.use_ball_neighbors_within_radius:
                idxs_ball = tree.query_ball_point(self.points[i], r=self.radius)
                if len(idxs_ball) == 0:
                    neigh_idx[out_i] = np.array([], dtype=int)
                    neigh_dist[out_i] = np.array([], dtype=float)
                    continue

                idxs_ball = np.array(idxs_ball, dtype=int)
                # Remove self if present
                mask_ball = (idxs_ball != i)
                idxs_ball = idxs_ball[mask_ball]
                if idxs_ball.size == 0:
                    neigh_idx[out_i] = np.array([], dtype=int)
                    neigh_dist[out_i] = np.array([], dtype=float)
                    continue

                dists_ball = np.linalg.norm(self.points[idxs_ball] - self.points[i], axis=1)
                # Order by (distance, index) deterministically
                order = np.lexsort((idxs_ball, dists_ball))
                neigh_idx[out_i] = idxs_ball[order]
                neigh_dist[out_i] = dists_ball[order]
            else:
                di, ii = tree.query(self.points[i], k=self.k + 1)
                ii = np.atleast_1d(ii)
                di = np.atleast_1d(di)
                mask = (ii != i)
                ii, di = ii[mask], di[mask]
                if ii.shape[0] < self.k:
                    K = min(self.k + 5, self.N - 1)
                    di2, ii2 = tree.query(self.points[i], k=K + 1)
                    m = (ii2 != i)
                    ii2, di2 = ii2[m], di2[m]
                    order = np.lexsort((ii2, di2))
                    row_idx, row_dist = ii2[order][: self.k], di2[order][: self.k]
                else:
                    order = np.lexsort((ii, di))
                    row_idx, row_dist = ii[order][: self.k], di[order][: self.k]

                neigh_idx[out_i] = row_idx
                neigh_dist[out_i] = row_dist

        return neigh_idx, neigh_dist

    def _estimate_normals(self, neigh_idx, indices=None):
        """
        Estimate normals with Open3D, then enforce deterministic sign
        relative to neighbor centroid vector.
        """
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(self.points)
        # Estimate normals. If using radius-mode neighbors, use radius search
        # so Open3D's estimation aligns with the neighbor selection. Otherwise
        # use KNN as before.
        if self.use_ball_neighbors_within_radius:
            pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamRadius(radius=self.radius))
        else:
            pcd.estimate_normals(search_param=o3d.geometry.KDTreeSearchParamKNN(knn=self.k))
        try:
            # Not relied upon for determinism, but can improve local smoothness
            pcd.orient_normals_consistent_tangent_plane(self.k)
        except Exception:
            pass

        normals = np.asarray(pcd.normals, dtype=np.float64)

        # Deterministic sign resolution: toward/away from neighbor centroid
        # We only resolve signs for the requested indices (or all points if
        # indices is None). `neigh_idx` is a list aligned with `indices` or
        # with all points when indices is None.
        if indices is None:
            idxs_to_process = np.arange(self.N)
        else:
            idxs_to_process = np.asarray(indices, dtype=int)

        # Build refs for selected points
        M = len(idxs_to_process)
        refs = np.zeros((M, 3), dtype=np.float64)
        for out_i, i in enumerate(idxs_to_process):
            idxs = neigh_idx[out_i]
            if idxs is None or len(idxs) == 0:
                refs[out_i] = np.zeros(3, dtype=np.float64)
            else:
                refs[out_i] = np.mean(self.points[np.asarray(idxs, dtype=int)], axis=0) - self.points[i]

        ref_norms = np.linalg.norm(refs, axis=1)
        valid = ref_norms > 0
        # Normals for the selected points
        normals_sel = normals[idxs_to_process]
        dots = np.sum(normals_sel * refs, axis=1)

        if self.normal_toward_centroid:
            # Ensure <n, ref> >= 0  (inward on typical watertight models)
            flip_mask = valid & (dots < 0)
        else:
            # Ensure <n, ref> <= 0  (outward on typical watertight models)
            flip_mask = valid & (dots > 0)

        if np.any(flip_mask):
            # Flip normals in the global normals array at the selected indices
            normals[idxs_to_process[flip_mask]] *= -1.0

        return normals

    def compute_lrf(self) -> np.ndarray:
        """
        Compute LRFs and return an (N, 3, 3) array of row-basis [x, y, z].
        """
        neigh_idx, _ = self._build_neighbors()
        # Determine which indices we're processing
        if self.indices is None:
            idxs_to_process = np.arange(self.N)
        else:
            idxs_to_process = self.indices

        normals = self._estimate_normals(neigh_idx, indices=idxs_to_process)

        # Prepare outputs
        lrf = np.zeros((self.N, 3, 3), dtype=np.float64)

        # Reference vectors for degeneracy fallback
        ref_vecs = np.array([[1.0, 0.0, 0.0],
                             [0.0, 1.0, 0.0],
                             [0.0, 0.0, 1.0]], dtype=np.float64)

        # Output shape: only for requested indices
        M = len(idxs_to_process)
        lrf_out = np.zeros((M, 3, 3), dtype=np.float64)

        for out_i, i in enumerate(idxs_to_process):
            p = self.points[i]
            n = normals[i]
            n_norm = np.linalg.norm(n)
            if n_norm == 0:
                # Fallback: reconstruct normal from local PCA if Open3D failed
                nbr = self.points[np.asarray(neigh_idx[out_i], dtype=int)] - p
                if nbr.shape[0] >= 3:
                    C = np.cov(nbr.T, bias=True)
                    eigvals, eigvecs = np.linalg.eigh(C)
                    n = eigvecs[:, np.argmin(eigvals)]
                    n_norm = np.linalg.norm(n)
                if n_norm == 0:
                    n = np.array([0, 0, 1.0])
                    n_norm = 1.0
            n = n / n_norm

            # Project neighbor offsets onto the tangent plane
            offsets = self.points[np.asarray(neigh_idx[out_i], dtype=int)] - p                 # (k,3)
            # Remove normal component: proj = v - (v·n) n
            projs = offsets - np.einsum('ij,j->i', offsets, n)[:, None] * n
            proj_norms = np.linalg.norm(projs, axis=1)

            # Keep only sufficiently non-degenerate projections
            keep = proj_norms > self.degeneracy_eps
            P = projs[keep]
            if P.shape[0] >= 2:
                # Center for PCA
                P_centered = P - np.mean(P, axis=0, keepdims=True)
                # SVD: P_centered = U S Vt, rows in R^3; principal dir is Vt[0]
                try:
                    U, S, Vt = np.linalg.svd(P_centered, full_matrices=False)
                    x = Vt[0]  # principal direction in R^3 (lies in tangent plane)
                except Exception:
                    # Fallback to covariance eigen-decomposition
                    C = np.cov(P_centered.T, bias=True)
                    eigvals, eigvecs = np.linalg.eigh(C)
                    x = eigvecs[:, np.argmax(eigvals)]
            else:
                # Severe degeneracy: use a projected reference axis
                x = None
                for rv in ref_vecs:
                    cand = rv - np.dot(rv, n) * n
                    if np.linalg.norm(cand) > self.degeneracy_eps:
                        x = cand
                        break
                if x is None:
                    # Last resort: pick any axis orthogonal to n
                    x = np.cross(n, np.array([1.0, 0.0, 0.0]))
                    if np.linalg.norm(x) < self.degeneracy_eps:
                        x = np.cross(n, np.array([0.0, 1.0, 0.0]))

            # Ensure x lies in tangent plane numerically
            x = x - np.dot(x, n) * n
            x_norm = np.linalg.norm(x)
            if x_norm == 0:
                x = np.cross(n, np.array([1.0, 0.0, 0.0]))
                if np.linalg.norm(x) < self.degeneracy_eps:
                    x = np.cross(n, np.array([0.0, 1.0, 0.0]))
                x_norm = np.linalg.norm(x)
            x = x / x_norm

            # Stable sign: align x with mean projected neighbor vector (fallback to strongest)
            if P.shape[0] > 0:
                mean_proj = P.sum(axis=0)
                mean_proj -= np.dot(mean_proj, n) * n
                if np.linalg.norm(mean_proj) <= self.degeneracy_eps:
                    # Fallback to the largest-magnitude projection among all neighbors
                        j = np.argmax(proj_norms)
                        mean_proj = projs[j]
                if np.dot(x, mean_proj) < 0:
                    x *= -1.0

            # y = z × x (right-handed)
            y = np.cross(n, x)
            y_norm = np.linalg.norm(y)
            if y_norm == 0:
                # Fallback if x ≈ n
                x = np.cross(np.array([1.0, 0.0, 0.0]), n)
                if np.linalg.norm(x) < self.degeneracy_eps:
                    x = np.cross(np.array([0.0, 1.0, 0.0]), n)
                x = x / (np.linalg.norm(x) or 1.0)
                y = np.cross(n, x)
                y = y / (np.linalg.norm(y) or 1.0)
            else:
                y = y / y_norm

            # Assemble preliminary frame
            M = np.vstack([x, y, n])

            # Per-point SVD re-orthonormalization to nearest rotation
            try:
                U, _, Vt = np.linalg.svd(M)
                R = U @ Vt
                # Ensure proper rotation (det > 0)
                if np.linalg.det(R) < 0:
                    R[0, :] *= -1.0
                lrf_out[out_i] = R
            except Exception:
                # Gram-Schmidt fallback
                x_gs = M[0]; x_gs = x_gs / (np.linalg.norm(x_gs) or 1.0)
                z_gs = M[2]; z_gs = z_gs / (np.linalg.norm(z_gs) or 1.0)
                y_gs = np.cross(z_gs, x_gs); y_gs = y_gs / (np.linalg.norm(y_gs) or 1.0)
                x_gs = np.cross(y_gs, z_gs)
                lrf_out[out_i] = np.vstack([x_gs, y_gs, z_gs])

        # Optional global convention: mean z >= 0 (does not flip normals themselves)
        if self.enforce_global_z_positive:
            mean_z = np.mean(lrf_out[:, 2, 2])
            if mean_z < 0:
                lrf_out[:, 0:2, :] *= -1.0

        self.lrf = lrf_out
        return self.lrf

    @classmethod
    def get_lrf(
        cls,
        points: np.ndarray,
        indices: np.ndarray = None,
        k: int = 20,
        radius: int = 20,
        use_ball_neighbors_within_radius: bool = False,
        normal_toward_centroid: bool = True,
        enforce_global_z_positive: bool = True,
        degeneracy_eps: float = 1e-8,
    ) -> np.ndarray:
        estimator = cls(
            points,
            indices=indices,
            k=k,
            radius=radius,
            use_ball_neighbors_within_radius=use_ball_neighbors_within_radius,
            normal_toward_centroid=normal_toward_centroid,
            enforce_global_z_positive=enforce_global_z_positive,
            degeneracy_eps=degeneracy_eps,
        )
        return estimator.compute_lrf()
