from __future__ import annotations

import numpy as np


def fit_plane_ransac(
    points: np.ndarray,
    distance_threshold: float = 0.3,
    ransac_n: int = 3,
    num_iterations: int = 100
) -> np.ndarray:
    """NumPy 기반 순수 수학 RANSAC 평면 피팅.
    
    포인트 클라우드에서 최적의 평면 방정식 ax + by + cz + d = 0을 구하고
    평면과의 거리가 임계값 이내인 인라이어(지면)의 불리언 마스크를 반환합니다.
    """
    n_points = len(points)
    if n_points < ransac_n:
        return np.ones(n_points, dtype=bool)

    best_inliers = np.zeros(n_points, dtype=bool)
    max_inlier_count = -1
    
    rng = np.random.default_rng(42)

    for _ in range(num_iterations):
        # 1. 무작위로 3개의 점 선택
        idx = rng.choice(n_points, ransac_n, replace=False)
        p0, p1, p2 = points[idx]

        # 2. 3개의 점으로 구성된 두 벡터
        v1 = p1 - p0
        v2 = p2 - p0

        # 3. 외적으로 평면의 법선 벡터(Normal Vector) 산출
        normal = np.cross(v1, v2)
        norm = np.linalg.norm(normal)
        if norm < 1e-6:
            continue  # 세 점이 일직선상에 있는 경우 스킵
        
        normal = normal / norm
        a, b, c = normal
        d = -np.dot(normal, p0)

        # 4. 평면 평정식과 모든 점 사이의 수직 거리 계산
        # distance = |ax + by + cz + d| / sqrt(a^2 + b^2 + c^2)
        # 법선 벡터가 이미 단위 벡터이므로 분모는 1
        distances = np.abs(np.dot(points, normal) + d)
        
        # 5. 거리가 threshold 내인 점들을 인라이어로 판정
        inliers = distances < distance_threshold
        inlier_count = np.sum(inliers)

        if inlier_count > max_inlier_count:
            max_inlier_count = inlier_count
            best_inliers = inliers

    # 평면 피팅이 정상 진행되었으면 인라이어 반환
    if max_inlier_count > 0:
        return best_inliers
    return np.ones(n_points, dtype=bool)


def filter_ransac(
    structured_data: dict | np.lib.npyio.NpzFile,
    distance_threshold: float = 0.3,
    ransac_n: int = 3,
    num_iterations: int = 100,
    local_grid_size: float = 10.0,
    progress_callback = None
) -> np.ndarray:
    """Local RANSAC 평면 분할 지면 필터 (Pure NumPy 버젼).
    
    외부 C++ 라이브러리(Open3D 등) 의존성 없이 산악 지형의 곡률에 대응하기 위해
    local_grid_size(기본 10m) 단위로 공간을 격자 분할한 뒤,
    각 격자별로 RANSAC 평면 피팅을 수행하여 지면 포인트(Inliers)를 취합합니다.
    """
    xyz = structured_data["xyz"]
    bounds = structured_data["bounds"]
    
    n_points = len(xyz)
    mask = np.zeros(n_points, dtype=bool)

    x_min, y_min = bounds[0][0], bounds[1][0]
    x_max, y_max = bounds[0][1], bounds[1][1]

    # 로컬 격자 구성
    grid_w = int(np.ceil((x_max - x_min) / local_grid_size)) + 1
    grid_h = int(np.ceil((y_max - y_min) / local_grid_size)) + 1

    xs = xyz[:, 0]
    ys = xyz[:, 1]

    gx = np.clip(((xs - x_min) / local_grid_size).astype(np.int32), 0, grid_w - 1)
    gy = np.clip(((ys - y_min) / local_grid_size).astype(np.int32), 0, grid_h - 1)

    grid_indices = gy * grid_w + gx
    unique_grids = np.unique(grid_indices)
    total_grids = len(unique_grids)

    # 격자 단위 피팅 실행
    for i, grid_id in enumerate(unique_grids):
        cell_mask = (grid_indices == grid_id)
        cell_points_idx = np.where(cell_mask)[0]
        
        if len(cell_points_idx) < ransac_n:
            mask[cell_points_idx] = True
            continue

        cell_xyz = xyz[cell_points_idx]
        
        # 순수 NumPy RANSAC 평면 피팅 실행
        cell_inliers = fit_plane_ransac(
            cell_xyz,
            distance_threshold=distance_threshold,
            ransac_n=ransac_n,
            num_iterations=num_iterations
        )
        
        # 인라이어 인덱스 취합
        global_inliers_idx = cell_points_idx[cell_inliers]
        mask[global_inliers_idx] = True

        # 진행도 한 줄 갱신 출력 대신 콜백 트리거
        if progress_callback:
            pct = int(((i + 1) / total_grids) * 100)
            progress_callback(pct)

    return mask