from __future__ import annotations

from dataclasses import dataclass

import numpy as np
from scipy.spatial import cKDTree

from .roi import apply_passthrough, normalize_aabb


@dataclass(slots=True)
class PreprocessResult:
    raw_points: np.ndarray
    filtered_points: np.ndarray
    work_points: np.ndarray
    work_normals: np.ndarray
    work_curvature: np.ndarray
    spacing: float
    stats: dict


def estimate_spacing(points: np.ndarray, sample_size: int = 4000) -> float:
    if len(points) < 2:
        return 1.0
    if len(points) > sample_size:
        indices = np.linspace(0, len(points) - 1, sample_size, dtype=int)
        sample = points[indices]
    else:
        sample = points
    tree = cKDTree(sample)
    distances, _ = tree.query(sample, k=min(2, len(sample)))
    if distances.ndim == 1:
        return float(np.median(distances))
    return float(np.quantile(distances[:, -1], 0.3))


def statistical_outlier_filter(
    points: np.ndarray,
    k: int,
    std_ratio: float,
) -> tuple[np.ndarray, np.ndarray]:
    if len(points) <= k + 1:
        mask = np.ones(len(points), dtype=bool)
        return points, mask
    tree = cKDTree(points)
    distances, _ = tree.query(points, k=k + 1)
    mean_dist = distances[:, 1:].mean(axis=1)
    threshold = mean_dist.mean() + std_ratio * mean_dist.std(ddof=0)
    mask = mean_dist <= threshold
    return points[mask], mask


def radius_filter(
    points: np.ndarray,
    radius: float,
    min_neighbors: int,
) -> tuple[np.ndarray, np.ndarray]:
    if len(points) == 0:
        return points, np.zeros(0, dtype=bool)
    tree = cKDTree(points)
    counts = tree.query_ball_point(points, r=radius, return_length=True)
    mask = counts >= min_neighbors
    return points[mask], mask


def voxel_downsample(points: np.ndarray, voxel_size: float) -> np.ndarray:
    if len(points) == 0 or voxel_size <= 0:
        return points.copy()
    mins = points.min(axis=0)
    grid = np.floor((points - mins) / voxel_size).astype(np.int64)
    unique, inverse = np.unique(grid, axis=0, return_inverse=True)
    sums = np.zeros((len(unique), 3), dtype=np.float64)
    counts = np.bincount(inverse)
    np.add.at(sums, inverse, points)
    return sums / counts[:, None]


def estimate_normals(
    points: np.ndarray,
    k: int,
) -> tuple[np.ndarray, np.ndarray]:
    if len(points) == 0:
        return np.zeros((0, 3)), np.zeros(0)
    neighbor_count = min(k + 1, len(points))
    tree = cKDTree(points)
    _, indices = tree.query(points, k=neighbor_count)
    normals = np.zeros((len(points), 3), dtype=np.float64)
    curvature = np.zeros(len(points), dtype=np.float64)
    for idx, neigh in enumerate(indices):
        neighborhood = points[np.atleast_1d(neigh)]
        centered = neighborhood - neighborhood.mean(axis=0, keepdims=True)
        cov = centered.T @ centered / max(len(centered), 1)
        eigenvalues, eigenvectors = np.linalg.eigh(cov)
        order = np.argsort(eigenvalues)
        eigenvalues = eigenvalues[order]
        eigenvectors = eigenvectors[:, order]
        normals[idx] = eigenvectors[:, 0]
        curvature[idx] = float(eigenvalues[0] / max(eigenvalues.sum(), 1e-9))
    return normals, curvature


def bilateral_filter_points(
    points: np.ndarray,
    normals: np.ndarray,
    *,
    k: int,
    sigma_d_mm: float,
    sigma_n_rad: float,
    iterations: int = 1,
) -> np.ndarray:
    """Point-cloud bilateral filter that smooths along estimated normals.

    Each point is shifted *along its own normal* by the weighted average of
    the signed normal-direction offsets to its k-nearest neighbours, using a
    Gaussian on Euclidean distance (``sigma_d_mm``) and a Gaussian on the
    angular deviation between point normals (``sigma_n_rad``). The tangential
    position is preserved, so edges survive while noise along the surface
    normal is suppressed.

    The per-point shift is clamped to ``sigma_d_mm`` so that degenerate
    neighbourhoods or bad normals cannot collapse the cloud.
    """
    if len(points) == 0 or iterations <= 0:
        return points.copy()
    effective_k = int(min(max(k, 2), len(points)))
    sigma_d = max(float(sigma_d_mm), 1e-6)
    sigma_n = max(float(sigma_n_rad), 1e-6)
    inv_2sd2 = 1.0 / (2.0 * sigma_d * sigma_d)
    inv_2sn2 = 1.0 / (2.0 * sigma_n * sigma_n)
    max_shift = sigma_d  # clamp: a bilateral step never moves a point further than sigma_d

    current = points.astype(np.float64, copy=True)
    unit_normals = normals.astype(np.float64, copy=True)
    norm_len = np.linalg.norm(unit_normals, axis=1)
    valid_normal = norm_len > 1e-9
    unit_normals[valid_normal] /= norm_len[valid_normal, None]

    for _ in range(int(iterations)):
        tree = cKDTree(current)
        distances, indices = tree.query(current, k=effective_k)
        if distances.ndim == 1:
            distances = distances[:, None]
            indices = indices[:, None]
        own_n = unit_normals[:, None, :]                       # (N, 1, 3)
        offsets = current[indices] - current[:, None, :]       # (N, k, 3)
        normal_offset = np.einsum("nkd,nkd->nk", offsets, own_n)
        spatial_w = np.exp(-(distances ** 2) * inv_2sd2)
        dots = np.clip(np.einsum("nkd,nkd->nk", unit_normals[indices], own_n), -1.0, 1.0)
        angle = np.arccos(np.abs(dots))
        normal_w = np.exp(-(angle ** 2) * inv_2sn2)
        weights = spatial_w * normal_w
        weight_sum = weights.sum(axis=1)
        move = weight_sum > 1e-12
        move &= valid_normal
        delta = np.zeros(len(current), dtype=np.float64)
        delta[move] = (weights[move] * normal_offset[move]).sum(axis=1) / weight_sum[move]
        delta = np.clip(delta, -max_shift, max_shift)
        current[move] = current[move] + delta[move, None] * unit_normals[move]
    return current


def preprocess_point_cloud(points: np.ndarray, config: dict) -> PreprocessResult:
    preprocess = config["preprocess"]
    flags: list[str] = []
    spacing_estimate = estimate_spacing(
        points,
        sample_size=preprocess["spacing_sample_size"],
    )

    passthrough_cfg = preprocess.get("passthrough") or {}
    passthrough_aabb = passthrough_cfg.get("aabb") if passthrough_cfg.get("enabled") else None
    raw_count = int(len(points))
    if passthrough_aabb is not None:
        aabb = normalize_aabb(passthrough_aabb)
        keep_mask = apply_passthrough(points, aabb)
        kept = points[keep_mask]
        if len(kept) >= max(64, raw_count // 1000):
            passthrough_points = kept
        else:
            passthrough_points = points
            flags.append("passthrough_empty")
    else:
        passthrough_points = points

    stat_points, stat_mask = statistical_outlier_filter(
        passthrough_points,
        k=preprocess["stat_k"],
        std_ratio=preprocess["stat_std_ratio"],
    )
    radius_points, radius_mask = radius_filter(
        stat_points,
        radius=preprocess["radius_mm"],
        min_neighbors=preprocess["radius_min_neighbors"],
    )

    bilateral_cfg = preprocess.get("bilateral") or {}
    bilateral_count = len(radius_points)
    bilateral_applied = False
    if bilateral_cfg.get("enabled") and len(radius_points) >= 8:
        k_requested = int(bilateral_cfg.get("k", 24))
        k_effective = min(k_requested, len(radius_points))
        if k_effective < k_requested:
            flags.append("bilateral_degraded")
        pre_normals, _ = estimate_normals(radius_points, k=min(k_effective, 32))
        radius_points = bilateral_filter_points(
            radius_points,
            pre_normals,
            k=k_effective,
            sigma_d_mm=float(bilateral_cfg.get("sigma_d_mm", max(2.0 * spacing_estimate, 2.0))),
            sigma_n_rad=float(
                bilateral_cfg.get(
                    "sigma_n_rad",
                    np.deg2rad(float(bilateral_cfg.get("sigma_n_deg", 15.0))),
                )
            ),
            iterations=int(bilateral_cfg.get("iterations", 1)),
        )
        bilateral_applied = True

    work_points = voxel_downsample(radius_points, preprocess["voxel_mm"])
    work_normals, work_curvature = estimate_normals(
        work_points,
        k=preprocess["normal_k"],
    )
    stats = {
        "raw_count": raw_count,
        "after_passthrough": int(len(passthrough_points)),
        "passthrough_applied": passthrough_aabb is not None,
        "after_statistical_filter": int(len(stat_points)),
        "after_radius_filter": int(len(radius_points)),
        "after_bilateral": int(bilateral_count),
        "bilateral_applied": bilateral_applied,
        "downsampled_count": int(len(work_points)),
        "spacing_mm": float(spacing_estimate),
        "statistical_keep_ratio": float(np.mean(stat_mask)) if len(stat_mask) else 1.0,
        "radius_keep_ratio": float(np.mean(radius_mask)) if len(radius_mask) else 1.0,
        "flags": flags,
    }
    return PreprocessResult(
        raw_points=points,
        filtered_points=radius_points,
        work_points=work_points,
        work_normals=work_normals,
        work_curvature=work_curvature,
        spacing=spacing_estimate,
        stats=stats,
    )