from __future__ import annotations

import numpy as np


def filter_csf(
    structured_data: dict | np.lib.npyio.NpzFile,
    cloth_resolution: float = 1.5,
    rigidness: int = 1,
    time_step: float = 0.65,
    class_threshold: float = 0.5,
    slope_smooth: bool = True
) -> np.ndarray:
    """Pure NumPy 기반 Cloth Simulation Filter (CSF) 보정본.
    
    물리 시뮬레이션의 반복 횟수를 충분히 확보(기본 150회)하여,
    가상의 천이 뒤집힌 임도 지형 최하단(원 지면의 최상단 봉우리)까지 확실하게 낙하하여 안착하도록 보장합니다.
    """
    xyz = structured_data["xyz"]
    n_points = len(xyz)
    if n_points == 0:
        return np.zeros(0, dtype=bool)

    xs = xyz[:, 0]
    ys = xyz[:, 1]
    zs = xyz[:, 2]
    
    # 1. 지형 반전 (Inversion)
    # 지표면 추출을 위해 높이를 뒤집습니다.
    z_max = np.max(zs)
    inverted_zs = z_max - zs

    # 2. 2D 가상 천 격자 설정 (바운더리 밀착 매핑)
    x_min, x_max = np.min(xs), np.max(xs)
    y_min, y_max = np.min(ys), np.max(ys)
    
    # 격자 경계 마진
    margin = cloth_resolution * 0.5
    cols = int(np.ceil((x_max - x_min) / cloth_resolution)) + 1
    rows = int(np.ceil((y_max - y_min) / cloth_resolution)) + 1
    
    # 천 노드의 Z 높이 초기화
    # 뒤집힌 지형의 최고 높이보다 약간 높은 곳에서 낙하 시작
    max_inverted_z = np.max(inverted_zs)
    start_height = max_inverted_z + 1.0
    cloth_z = np.full((rows, cols), start_height, dtype=np.float32)

    # 3. 격자 충돌 타겟 구성 (Drape Target)
    collision_grid = np.full((rows, cols), -np.inf, dtype=np.float32)
    
    # 포인트들의 격자 인덱스 투영 (소수점 탈락을 정교하게 제어)
    gx = np.clip(((xs - x_min) / cloth_resolution).astype(np.int32), 0, cols - 1)
    gy = np.clip(((ys - y_min) / cloth_resolution).astype(np.int32), 0, rows - 1)
    
    # 각 격자별 최댓값(뒤집힌 지면) 누적
    np.maximum.at(collision_grid, (gy, gx), inverted_zs.astype(np.float32))
    
    # 비어있는 격자는 안전 최저 높이(0.0)로 채움
    collision_grid[collision_grid == -np.inf] = 0.0

    # 4. 천 시뮬레이션 반복 루프 (물리 하강)
    # 충분히 낙하하도록 루프 횟수를 150회로 대폭 확장
    iterations = 150
    gravity = 9.8 * time_step * 0.05  # 중력 하강 속도
    
    # rigidness에 따른 완화 계수 설정
    # 1: 인장력이 연하여 지형을 바짝 밀착 (산악 지형)
    # 2: 중간 완화
    # 3: 단단하여 지형을 부드럽게 덮음
    spring_coeff = 0.25 if rigidness == 1 else (0.45 if rigidness == 2 else 0.65)

    for _ in range(iterations):
        # (1) 중력 강하
        cloth_z -= gravity
        
        # (2) 지형 충돌 검사 (Drape 충돌 하한선 제약)
        cloth_z = np.maximum(cloth_z, collision_grid)

        # (3) 노드 간 스프링 제약 조건 완화 (인접 픽셀 교정)
        # 가로 방향
        diff_h = cloth_z[:, 1:] - cloth_z[:, :-1]
        correction_h = diff_h * spring_coeff * 0.5
        cloth_z[:, :-1] += correction_h
        cloth_z[:, 1:] -= correction_h
        
        # 세로 방향
        diff_v = cloth_z[1:, :] - cloth_z[:-1, :]
        correction_v = diff_v * spring_coeff * 0.5
        cloth_z[:-1, :] += correction_v
        cloth_z[1:, :] -= correction_v

        # 제약 조건 완화 후 다시 지형 높이 준수 강제
        cloth_z = np.maximum(cloth_z, collision_grid)

    # 5. 시뮬레이션된 가상 천 높이와 원본 Z 높이 대조
    simulated_inverted_z = cloth_z[gy, gx]
    
    # 반전된 지형 오차 거리 산출
    height_diff = np.abs(inverted_zs - simulated_inverted_z)

    # 오차 거리가 threshold 이내에 들어온 포인트를 지면으로 분류
    mask = height_diff <= class_threshold

    # 수목 노이즈 2차 필터 보정
    if slope_smooth:
        local_min_z = collision_grid[gy, gx]
        mask = mask & ((inverted_zs - local_min_z) < 1.8)

    return mask