"""포인트클라우드 자체 분류 유틸리티.

원본 LAS에는 classification 라벨이 없으므로(전부 0), 데이터에 들어있는 신호로
우리가 직접 클래스를 부여한다. 두 가지 독립적인 분류 기준을 산출한다.

  1) RGB 기준  : 식생지수 ExG = 2g - r - b (정규화) 로 녹색식생 여부 판정
  2) 지면 추출 : 격자 최저점(grid min-Z) 표면 대비 높이로 지면 여부 판정

설계 의도
---------
- 무거운 작업(원본 전체 스캔 → 5백만 점 샘플 + 피처 계산)은 **한 번만** 수행해
  npz 피처 캐시로 저장한다.
- 사용자가 민감도(임계값)를 바꿀 때는 캐시된 피처에 임계값만 다시 적용하므로 빠르다.
- 분류 알고리즘을 교체/개선하려면 `classify_from_features` 와 각 Classifier 만
  수정하면 되고, 캐시 포맷이나 API 는 건드릴 필요가 없다.
"""

from __future__ import annotations

from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any

import laspy
import numpy as np


# ---- 파라미터 ----

# 클래스 코드 (두 기준 각각 독립)
RGB_NON_VEGETATION = 0   # 비식생 (지면/구조물/흙 등)
RGB_VEGETATION = 1       # 녹색식생

GROUND_NON_GROUND = 0    # 비지면 (지면 위 물체: 나무/구조물)
GROUND_GROUND = 1        # 지면

FEATURE_VERSION = 2      # 캐시 무효화용 버전


@dataclass
class ClassifyParams:
    """사용자 조절 가능한 민감도 파라미터.

    rgb_exg_threshold     : ExG 가 이 값보다 크면 식생. 낮출수록 식생을 민감하게 잡음.
    ground_height_threshold: 격자 최저점 대비 높이가 이 값 이하이면 지면.
                             높일수록 지면을 후하게(민감하지 않게) 잡음.
    """

    rgb_exg_threshold: float = 0.05
    ground_height_threshold: float = 1.5

    def clamped(self) -> "ClassifyParams":
        return ClassifyParams(
            rgb_exg_threshold=float(np.clip(self.rgb_exg_threshold, -0.5, 1.0)),
            ground_height_threshold=float(np.clip(self.ground_height_threshold, 0.1, 50.0)),
        )


# 피처 캐시 빌드용 고정 파라미터 (분류 임계값과 무관)
_DEFAULT_MAX_POINTS = 5_000_000
_DEFAULT_CELL_SIZE = 2.0


# ---- 피처 추출 (무거운 1회 작업) ----

def _normalize_rgb_channel(arr: np.ndarray) -> np.ndarray:
    """LAS RGB(uint16) 를 0..255 로 정규화. 16bit(>255)면 8bit 로 축소."""
    a = np.asarray(arr).astype(np.float32)
    if a.size and a.max() > 255.0:
        a = a / 256.0
    return np.clip(a, 0, 255)


def _excess_green(r: np.ndarray, g: np.ndarray, b: np.ndarray) -> np.ndarray:
    """정규화 ExG = 2g - r - b  (각 채널을 r+g+b 로 정규화한 뒤 계산)."""
    s = r + g + b
    s = np.where(s <= 0, 1.0, s)
    rn, gn, bn = r / s, g / s, b / s
    return (2.0 * gn - rn - bn).astype(np.float32)


def build_feature_cache(
    las_path: Path,
    cache_path: Path,
    max_points: int = _DEFAULT_MAX_POINTS,
    cell_size: float = _DEFAULT_CELL_SIZE,
) -> dict[str, Any]:
    """원본 LAS 를 스캔해 5백만 점 샘플 + 피처(ExG, 지면대비높이)를 npz 로 저장."""
    cache_path.parent.mkdir(parents=True, exist_ok=True)

    # 헤더에서 공간 범위 획득
    with laspy.open(las_path) as f:
        header = f.header
        total_points = int(header.point_count)
        x_min, y_min = float(header.mins[0]), float(header.mins[1])
        x_max, y_max = float(header.maxs[0]), float(header.maxs[1])
        z_min, z_max = float(header.mins[2]), float(header.maxs[2])
        point_format = header.point_format
        has_rgb = all(name in point_format.dimension_names for name in ("red", "green", "blue"))

    cs = cell_size
    grid_w = int(np.ceil((x_max - x_min) / cs)) + 2
    grid_h = int(np.ceil((y_max - y_min) / cs)) + 2
    min_z_grid = np.full((grid_h, grid_w), np.inf, dtype=np.float32)

    # Pass 1: 격자별 최저 Z (지면 표면 추정) — 전체 점 사용
    with laspy.open(las_path) as f:
        for chunk in f.chunk_iterator(500_000):
            xs = np.asarray(chunk.x, dtype=np.float32)
            ys = np.asarray(chunk.y, dtype=np.float32)
            zs = np.asarray(chunk.z, dtype=np.float32)
            gx = np.clip(((xs - x_min) / cs).astype(np.int32), 0, grid_w - 1)
            gy = np.clip(((ys - y_min) / cs).astype(np.int32), 0, grid_h - 1)
            np.minimum.at(min_z_grid, (gy, gx), zs)

    min_z_grid[min_z_grid == np.inf] = z_min
    try:
        from scipy.ndimage import minimum_filter as _mf
        min_z_grid = _mf(min_z_grid, size=3).astype(np.float32)
    except ImportError:
        pass

    # Pass 2: 균일 랜덤 샘플 + 피처 계산
    keep_prob = 1.0 if total_points <= max_points else max_points / total_points
    rng = np.random.default_rng(42)

    xyz_parts: list[np.ndarray] = []
    rgb_parts: list[np.ndarray] = []
    exg_parts: list[np.ndarray] = []
    hgt_parts: list[np.ndarray] = []
    collected = 0

    with laspy.open(las_path) as f:
        for chunk in f.chunk_iterator(500_000):
            n = len(chunk.x)
            mask = rng.random(n) < keep_prob if keep_prob < 1.0 else np.ones(n, dtype=bool)
            if not mask.any():
                continue
            xs = np.asarray(chunk.x, dtype=np.float64)[mask]
            ys = np.asarray(chunk.y, dtype=np.float64)[mask]
            zs = np.asarray(chunk.z, dtype=np.float64)[mask]

            if has_rgb:
                r = _normalize_rgb_channel(np.asarray(chunk.red)[mask])
                g = _normalize_rgb_channel(np.asarray(chunk.green)[mask])
                b = _normalize_rgb_channel(np.asarray(chunk.blue)[mask])
            else:
                r = g = b = np.full(int(mask.sum()), 128.0, dtype=np.float32)
            exg = _excess_green(r, g, b)

            gx = np.clip(((xs - x_min) / cs).astype(np.int32), 0, grid_w - 1)
            gy = np.clip(((ys - y_min) / cs).astype(np.int32), 0, grid_h - 1)
            hgt = (zs.astype(np.float32) - min_z_grid[gy, gx]).astype(np.float32)

            xyz_parts.append(np.stack([xs, ys, zs], axis=1).astype(np.float32))
            rgb_parts.append(np.stack([r, g, b], axis=1).astype(np.uint8))
            exg_parts.append(exg)
            hgt_parts.append(hgt)
            collected += int(mask.sum())

    if xyz_parts:
        xyz = np.concatenate(xyz_parts)
        rgb = np.concatenate(rgb_parts)
        exg = np.concatenate(exg_parts)
        hgt = np.concatenate(hgt_parts)
        if len(xyz) > max_points:
            sel = rng.choice(len(xyz), max_points, replace=False)
            xyz, rgb, exg, hgt = xyz[sel], rgb[sel], exg[sel], hgt[sel]
    else:
        xyz = np.empty((0, 3), np.float32)
        rgb = np.empty((0, 3), np.uint8)
        exg = np.empty((0,), np.float32)
        hgt = np.empty((0,), np.float32)

    np.savez_compressed(
        cache_path,
        version=np.array([FEATURE_VERSION], dtype=np.int32),
        xyz=xyz,
        rgb=rgb,
        exg=exg,
        hgt=hgt,
        bounds=np.array([[x_min, x_max], [y_min, y_max], [z_min, z_max]], dtype=np.float64),
        total_points=np.array([total_points], dtype=np.int64),
        has_rgb=np.array([1 if has_rgb else 0], dtype=np.int32),
        cell_size=np.array([cs], dtype=np.float32),
    )

    return {
        "sample_point_count": int(len(xyz)),
        "source_point_count": total_points,
        "has_rgb": has_rgb,
        "cell_size": cs,
    }


def load_feature_cache(cache_path: Path) -> dict[str, Any] | None:
    if not cache_path.exists():
        return None
    data = np.load(cache_path)
    if int(data["version"][0]) != FEATURE_VERSION:
        return None
    return {
        "xyz": data["xyz"],
        "rgb": data["rgb"],
        "exg": data["exg"],
        "hgt": data["hgt"],
        "bounds": data["bounds"],
        "total_points": int(data["total_points"][0]),
        "has_rgb": bool(data["has_rgb"][0]),
        "cell_size": float(data["cell_size"][0]),
    }


# ---- 분류 (가벼운 작업, 임계값만 적용) ----

def classify_from_features(
    features: dict[str, Any], params: ClassifyParams
) -> tuple[np.ndarray, np.ndarray]:
    """캐시된 피처에 민감도 임계값을 적용해 두 클래스 배열을 반환.

    returns (rgb_class, ground_class)  각 dtype=uint8, 길이=점 개수
    """
    p = params.clamped()
    exg = features["exg"]
    hgt = features["hgt"]

    rgb_class = np.where(exg > p.rgb_exg_threshold, RGB_VEGETATION, RGB_NON_VEGETATION).astype(np.uint8)
    ground_class = np.where(
        (hgt >= 0.0) & (hgt <= p.ground_height_threshold), GROUND_GROUND, GROUND_NON_GROUND
    ).astype(np.uint8)
    return rgb_class, ground_class


def run_classification(
    las_path: Path,
    cache_path: Path,
    params: ClassifyParams,
    max_points: int = _DEFAULT_MAX_POINTS,
) -> dict[str, Any]:
    """피처 캐시를 보장(없으면 생성)한 뒤 분류 결과를 프론트엔드용 dict 로 반환."""
    features = load_feature_cache(cache_path)
    if features is None:
        build_feature_cache(las_path, cache_path, max_points=max_points)
        features = load_feature_cache(cache_path)
    assert features is not None

    rgb_class, ground_class = classify_from_features(features, params)
    xyz = features["xyz"]
    bounds = features["bounds"]
    p = params.clamped()

    n = int(len(xyz))
    veg_count = int(np.count_nonzero(rgb_class == RGB_VEGETATION))
    ground_count = int(np.count_nonzero(ground_class == GROUND_GROUND))

    return {
        "source_point_count": features["total_points"],
        "sample_point_count": n,
        "has_rgb": features["has_rgb"],
        "browser_limit_applied": features["total_points"] > n,
        "bounds": {
            "x": [float(bounds[0][0]), float(bounds[0][1])],
            "y": [float(bounds[1][0]), float(bounds[1][1])],
            "z": [float(bounds[2][0]), float(bounds[2][1])],
        },
        "params": asdict(p),
        "points": np.round(xyz, 3).tolist(),
        "rgb": features["rgb"].tolist(),
        "rgb_class": rgb_class.tolist(),
        "ground_class": ground_class.tolist(),
        "summary": {
            "total_sample": n,
            "rgb": {
                "vegetation": veg_count,
                "non_vegetation": n - veg_count,
                "vegetation_pct": round(veg_count / n * 100, 1) if n else 0.0,
            },
            "ground": {
                "ground": ground_count,
                "non_ground": n - ground_count,
                "ground_pct": round(ground_count / n * 100, 1) if n else 0.0,
            },
        },
    }