"""直线拟合模块 - 使用RANSAC算法拟合3D直线"""
import numpy as np
from dataclasses import dataclass
from typing import Tuple, Optional


@dataclass
class LineModel:
    """直线模型"""
    point: np.ndarray      # 直线上的点 (3,)
    direction: np.ndarray  # 方向向量（单位向量） (3,)
    inliers: np.ndarray    # 内点索引
    inlier_ratio: float    # 内点比例

    def distance_to_point(self, point: np.ndarray) -> float:
        """计算点到直线的距离"""
        vector = point - self.point
        projection = np.dot(vector, self.direction) * self.direction
        perpendicular = vector - projection
        return np.linalg.norm(perpendicular)

    def distance_to_points(self, points: np.ndarray) -> np.ndarray:
        """计算多个点到直线的距离"""
        vectors = points - self.point
        projections = np.dot(vectors, self.direction)[:, np.newaxis] * self.direction
        perpendicular_vectors = vectors - projections
        distances = np.linalg.norm(perpendicular_vectors, axis=1)
        return distances

    def sample_points_on_line(self, n_samples: int, point_cloud: np.ndarray) -> np.ndarray:
        """
        从拟合的直线对应的真实点云中采样点

        Args:
            n_samples: 采样点数
            point_cloud: 用于拟合的原始点云（内点）

        Returns:
            采样的点 (n_samples, 3)
        """
        if len(point_cloud) == 0:
            return np.array([]).reshape(0, 3)

        n_samples = min(n_samples, len(point_cloud))
        indices = np.random.choice(len(point_cloud), size=n_samples, replace=False)
        return point_cloud[indices]


class LineFitter:
    """RANSAC直线拟合器"""

    def __init__(
        self,
        threshold: float = 2.0,
        max_iterations: int = 1000,
        min_inlier_ratio: float = 0.6
    ):
        """
        Args:
            threshold: 内点距离阈值(mm)
            max_iterations: 最大迭代次数
            min_inlier_ratio: 最小内点比例
        """
        self.threshold = threshold
        self.max_iterations = max_iterations
        self.min_inlier_ratio = min_inlier_ratio

    def fit(self, points: np.ndarray) -> Optional[LineModel]:
        """
        使用RANSAC拟合3D直线

        Args:
            points: 点云数组 (N, 3)

        Returns:
            LineModel或None（拟合失败）
        """
        if len(points) < 2:
            return None

        best_inliers = None
        best_point = None
        best_direction = None
        best_inlier_count = 0

        for _ in range(self.max_iterations):
            # 随机选择2个点
            sample_indices = np.random.choice(len(points), 2, replace=False)
            p1, p2 = points[sample_indices]

            # 计算直线方向
            direction = p2 - p1
            if np.linalg.norm(direction) < 1e-6:
                continue
            direction = direction / np.linalg.norm(direction)

            # 计算所有点到直线的距离
            vectors = points - p1
            projections = np.dot(vectors, direction)[:, np.newaxis] * direction
            distances = np.linalg.norm(vectors - projections, axis=1)

            # 找到内点
            inliers = np.where(distances < self.threshold)[0]

            if len(inliers) > best_inlier_count:
                best_inlier_count = len(inliers)
                best_inliers = inliers
                best_point = p1
                best_direction = direction

        if best_inliers is None or len(best_inliers) < 2:
            return None

        # 使用所有内点重新拟合直线（最小二乘）
        inlier_points = points[best_inliers]
        final_point, final_direction = self._fit_line_least_squares(inlier_points)

        if final_direction is None:
            return None

        # 重新计算内点
        vectors = points - final_point
        projections = np.dot(vectors, final_direction)[:, np.newaxis] * final_direction
        distances = np.linalg.norm(vectors - projections, axis=1)
        final_inliers = np.where(distances < self.threshold)[0]

        inlier_ratio = len(final_inliers) / len(points)

        # 检查内点比例
        if inlier_ratio < self.min_inlier_ratio:
            return None

        return LineModel(
            point=final_point,
            direction=final_direction,
            inliers=final_inliers,
            inlier_ratio=inlier_ratio
        )

    def _fit_line_least_squares(self, points: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        使用最小二乘法拟合直线

        Returns:
            (point, direction): 直线上的点和方向向量
        """
        if len(points) < 2:
            return None, None

        # 计算质心
        centroid = np.mean(points, axis=0)

        # 中心化点云
        centered = points - centroid

        # PCA：找主方向
        try:
            # 协方差矩阵
            cov = np.dot(centered.T, centered) / len(points)
            # 特征值分解
            eigenvalues, eigenvectors = np.linalg.eig(cov)
            # 最大特征值对应的特征向量就是直线方向
            max_idx = np.argmax(eigenvalues)
            direction = eigenvectors[:, max_idx].real
            direction = direction / np.linalg.norm(direction)

            return centroid, direction
        except np.linalg.LinAlgError:
            return None, None