#!/usr/bin/env python3
"""Feature-aware, topology-preserving mesh denoising experiments.

The implementation follows the two-stage family used by guided mesh-normal
filtering: robustly filter face normals, then move the *existing* vertices
toward the planes described by those normals.  It adds safeguards needed for
this product asset: a hard displacement trust region, feature/boundary locks,
mostly-normal motion, data fidelity, and reflection-field diagnostics.
"""

from __future__ import annotations

import argparse
import gc
import json
import math
import time
from dataclasses import asdict, dataclass
from pathlib import Path

import numpy as np


EPS = np.float32(1e-20)


@dataclass(frozen=True)
class Config:
    name: str
    guide_iterations: int
    normal_iterations: int
    normal_sigma: float
    vertex_iterations: int
    vertex_step: float
    data_pull: float
    max_displacement_diagonal: float
    tangent_ratio: float


CONFIGS = (
    Config("gentle", 2, 6, 0.30, 5, 0.52, 0.045, 0.00075, 0.06),
    Config("balanced", 3, 14, 0.36, 10, 0.62, 0.028, 0.00150, 0.08),
    Config("strong", 4, 28, 0.44, 16, 0.70, 0.018, 0.00300, 0.10),
    Config("wide", 10, 30, 0.52, 28, 0.82, 0.006, 0.00600, 0.12),
    Config("deep", 18, 50, 0.62, 46, 0.88, 0.004, 0.01000, 0.14),
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True)
    parser.add_argument("--outdir", required=True)
    parser.add_argument(
        "--variants",
        nargs="+",
        default=[config.name for config in CONFIGS],
        choices=[config.name for config in CONFIGS],
    )
    return parser.parse_args()


def normalize_rows(values: np.ndarray) -> np.ndarray:
    lengths = np.linalg.norm(values, axis=1)
    return values / np.maximum(lengths[:, None], EPS)


def face_geometry(
    vertices: np.ndarray, faces: np.ndarray, *, need_centroids: bool = True
) -> tuple[np.ndarray, np.ndarray, np.ndarray | None]:
    edge_a = vertices[faces[:, 1]] - vertices[faces[:, 0]]
    edge_b = vertices[faces[:, 2]] - vertices[faces[:, 0]]
    crosses = np.cross(edge_a, edge_b)
    double_area = np.linalg.norm(crosses, axis=1)
    normals = crosses / np.maximum(double_area[:, None], EPS)
    centroids = None
    if need_centroids:
        centroids = (
            vertices[faces[:, 0]]
            + vertices[faces[:, 1]]
            + vertices[faces[:, 2]]
        ) / np.float32(3.0)
    return normals.astype(np.float32), (double_area * 0.5).astype(np.float32), centroids


def vertex_normals(
    vertices: np.ndarray, faces: np.ndarray, face_areas: np.ndarray, face_normals: np.ndarray
) -> np.ndarray:
    weighted = face_normals * face_areas[:, None]
    result = np.zeros_like(vertices, dtype=np.float32)
    for corner in range(3):
        indices = faces[:, corner]
        for axis in range(3):
            result[:, axis] += np.bincount(
                indices,
                weights=weighted[:, axis],
                minlength=len(vertices),
            ).astype(np.float32)
    return normalize_rows(result).astype(np.float32)


def build_topology(faces: np.ndarray, vertex_count: int) -> dict[str, np.ndarray | int]:
    """Build edge-adjacent face pairs without Python dictionaries."""
    face_count = len(faces)
    raw_edges = np.empty((face_count * 3, 2), dtype=np.int32)
    raw_edges[0::3] = faces[:, (0, 1)]
    raw_edges[1::3] = faces[:, (1, 2)]
    raw_edges[2::3] = faces[:, (2, 0)]
    raw_edges.sort(axis=1)
    owners = np.repeat(np.arange(face_count, dtype=np.int32), 3)

    keys = (raw_edges[:, 0].astype(np.uint64) << np.uint64(32)) | raw_edges[:, 1].astype(
        np.uint64
    )
    order = np.argsort(keys, kind="stable")
    sorted_keys = keys[order]
    group_starts = np.flatnonzero(
        np.r_[True, sorted_keys[1:] != sorted_keys[:-1]]
    ).astype(np.int64)
    group_ends = np.r_[group_starts[1:], len(sorted_keys)]
    group_counts = group_ends - group_starts

    unique_indices = order[group_starts]
    unique_edges = raw_edges[unique_indices].copy()
    manifold = group_counts == 2
    starts = group_starts[manifold]
    raw_first = order[starts]
    raw_second = order[starts + 1]
    adjacent_faces = np.column_stack((owners[raw_first], owners[raw_second])).astype(np.int32)
    adjacent_edges = raw_edges[raw_first].copy()

    boundary_groups = group_counts == 1
    boundary_edges = raw_edges[order[group_starts[boundary_groups]]].copy()
    boundary_vertices = np.zeros(vertex_count, dtype=bool)
    if len(boundary_edges):
        boundary_vertices[boundary_edges.ravel()] = True

    stats = {
        "unique_edge_count": int(len(unique_edges)),
        "boundary_edge_count": int(np.count_nonzero(boundary_groups)),
        "nonmanifold_edge_count": int(np.count_nonzero(group_counts > 2)),
    }
    del raw_edges, owners, keys, order, sorted_keys, group_starts, group_ends, group_counts
    gc.collect()
    return {
        "unique_edges": unique_edges,
        "adjacent_faces": adjacent_faces,
        "adjacent_edges": adjacent_edges,
        "boundary_vertices": boundary_vertices,
        "stats": stats,
    }


def mean_edge_length(vertices: np.ndarray, edges: np.ndarray) -> float:
    total = 0.0
    chunk = 500_000
    for start in range(0, len(edges), chunk):
        part = edges[start : start + chunk]
        delta = vertices[part[:, 0]] - vertices[part[:, 1]]
        total += float(np.linalg.norm(delta, axis=1).sum(dtype=np.float64))
    return total / max(len(edges), 1)


def diffuse_normals(
    normals: np.ndarray,
    face_areas: np.ndarray,
    adjacent_faces: np.ndarray,
    iterations: int,
    sigma: float,
    *,
    fixed_guidance: np.ndarray | None = None,
    blend: float = 1.0,
) -> np.ndarray:
    """Area-weighted bilateral diffusion over edge-adjacent faces."""
    left = adjacent_faces[:, 0]
    right = adjacent_faces[:, 1]
    current = normals.copy()
    area_scale = np.maximum(face_areas / np.median(face_areas), np.float32(1e-4))
    area_scale = np.minimum(area_scale, np.float32(8.0)).astype(np.float32)
    for _ in range(iterations):
        guidance = current if fixed_guidance is None else fixed_guidance
        dots = np.einsum("ij,ij->i", guidance[left], guidance[right])
        normal_distance = np.sqrt(np.maximum(0.0, 2.0 - 2.0 * np.clip(dots, -1.0, 1.0)))
        pair_weight = np.exp(-0.5 * np.square(normal_distance / sigma)).astype(np.float32)

        accumulated = current * area_scale[:, None]
        weight_sum = area_scale.copy()
        for destination, source in ((left, right), (right, left)):
            weights = pair_weight * area_scale[source]
            weight_sum += np.bincount(
                destination, weights=weights, minlength=len(current)
            ).astype(np.float32)
            for axis in range(3):
                accumulated[:, axis] += np.bincount(
                    destination,
                    weights=weights * current[source, axis],
                    minlength=len(current),
                ).astype(np.float32)
        updated = normalize_rows(accumulated / np.maximum(weight_sum[:, None], EPS))
        current = normalize_rows((1.0 - blend) * current + blend * updated).astype(np.float32)
    return current


def smoothstep(edge0: float, edge1: float, values: np.ndarray) -> np.ndarray:
    scaled = np.clip((values - edge0) / (edge1 - edge0), 0.0, 1.0)
    return scaled * scaled * (3.0 - 2.0 * scaled)


def feature_protection(
    guidance: np.ndarray,
    adjacent_faces: np.ndarray,
    adjacent_edges: np.ndarray,
    boundary_vertices: np.ndarray,
    vertex_count: int,
) -> tuple[np.ndarray, np.ndarray]:
    left, right = adjacent_faces[:, 0], adjacent_faces[:, 1]
    dots = np.einsum("ij,ij->i", guidance[left], guidance[right])
    distance = np.sqrt(np.maximum(0.0, 2.0 - 2.0 * np.clip(dots, -1.0, 1.0)))
    # About 12.6 degrees begins protection; 32 degrees is fully locked.
    edge_strength = smoothstep(0.22, 0.55, distance).astype(np.float32)
    vertex_strength = np.zeros(vertex_count, dtype=np.float32)
    np.maximum.at(vertex_strength, adjacent_edges[:, 0], edge_strength)
    np.maximum.at(vertex_strength, adjacent_edges[:, 1], edge_strength)
    vertex_strength[boundary_vertices] = 1.0
    return vertex_strength, distance.astype(np.float32)


def update_vertices(
    original_vertices: np.ndarray,
    faces: np.ndarray,
    target_normals: np.ndarray,
    original_vertex_normals: np.ndarray,
    feature_strength: np.ndarray,
    config: Config,
    diagonal: float,
) -> np.ndarray:
    vertices = original_vertices.copy()
    counts = np.zeros(len(vertices), dtype=np.float32)
    for corner in range(3):
        counts += np.bincount(faces[:, corner], minlength=len(vertices)).astype(np.float32)
    counts = np.maximum(counts, 1.0)
    trust_radius = np.float32(diagonal * config.max_displacement_diagonal)

    for _ in range(config.vertex_iterations):
        centroids = (
            vertices[faces[:, 0]] + vertices[faces[:, 1]] + vertices[faces[:, 2]]
        ) / np.float32(3.0)
        correction = np.zeros_like(vertices, dtype=np.float32)
        for corner in range(3):
            indices = faces[:, corner]
            residual = np.einsum("ij,ij->i", target_normals, centroids - vertices[indices])
            projected = target_normals * residual[:, None]
            for axis in range(3):
                correction[:, axis] += np.bincount(
                    indices,
                    weights=projected[:, axis],
                    minlength=len(vertices),
                ).astype(np.float32)
        correction /= counts[:, None]
        correction *= (1.0 - 0.92 * feature_strength[:, None])
        proposed = vertices + np.float32(config.vertex_step) * correction

        displacement = proposed - original_vertices
        normal_amount = np.einsum("ij,ij->i", displacement, original_vertex_normals)
        normal_part = original_vertex_normals * normal_amount[:, None]
        tangent_part = displacement - normal_part
        displacement = normal_part + np.float32(config.tangent_ratio) * tangent_part
        displacement *= np.float32(1.0 - config.data_pull)

        lengths = np.linalg.norm(displacement, axis=1)
        scale = np.minimum(1.0, trust_radius / np.maximum(lengths, EPS)).astype(np.float32)
        displacement *= scale[:, None]
        # Prevent a tiny numerical translation from masquerading as surface change.
        displacement -= displacement.mean(axis=0, dtype=np.float64).astype(np.float32)
        vertices = original_vertices + displacement
    return vertices.astype(np.float32)


def revert_unsafe_faces(
    original_vertices: np.ndarray,
    candidate_vertices: np.ndarray,
    faces: np.ndarray,
    original_face_normals: np.ndarray,
    *,
    max_passes: int = 6,
) -> tuple[np.ndarray, int]:
    """Restore vertices incident to inverted/degenerate faces.

    The inputs can contain a few extremely small triangles where even a tiny
    trusted displacement changes orientation.  Reverting their incident
    vertices is conservative and affects a negligible fraction of the mesh.
    """
    vertices = candidate_vertices.copy()
    reverted = np.zeros(len(vertices), dtype=bool)
    for _ in range(max_passes):
        normals, areas, _ = face_geometry(vertices, faces, need_centroids=False)
        alignment = np.einsum("ij,ij->i", original_face_normals, normals)
        unsafe = (alignment <= 0.0) | (areas <= 1e-14)
        if not np.any(unsafe):
            break
        indices = np.unique(faces[unsafe].ravel())
        reverted[indices] = True
        vertices[indices] = original_vertices[indices]
    return vertices, int(np.count_nonzero(reverted))


def surface_area(vertices: np.ndarray, faces: np.ndarray) -> float:
    _, areas, _ = face_geometry(vertices, faces, need_centroids=False)
    return float(areas.sum(dtype=np.float64))


def signed_volume(vertices: np.ndarray, faces: np.ndarray) -> float:
    total = 0.0
    chunk = 400_000
    for start in range(0, len(faces), chunk):
        tri = faces[start : start + chunk]
        cross = np.cross(vertices[tri[:, 1]], vertices[tri[:, 2]])
        total += float(np.einsum("ij,ij->", vertices[tri[:, 0]], cross, dtype=np.float64))
    return total / 6.0


def reflection_rms(
    normals: np.ndarray, adjacent_faces: np.ndarray, smooth_mask: np.ndarray
) -> float:
    left = adjacent_faces[:, 0]
    right = adjacent_faces[:, 1]
    views = np.asarray(
        (
            (0.0, 0.0, 1.0),
            (0.0, -0.8944272, 0.4472136),
            (0.8728716, -0.4364358, 0.2182179),
            (-0.6666667, -0.6666667, 0.3333333),
        ),
        dtype=np.float32,
    )
    values: list[float] = []
    for view in views:
        dots = normals @ view
        reflected = 2.0 * dots[:, None] * normals - view[None, :]
        delta = reflected[left] - reflected[right]
        squared = np.einsum("ij,ij->i", delta, delta)
        values.append(float(np.sqrt(np.mean(squared[smooth_mask], dtype=np.float64))))
    return float(np.mean(values))


def geometry_metrics(
    original_vertices: np.ndarray,
    vertices: np.ndarray,
    faces: np.ndarray,
    original_normals: np.ndarray,
    adjacent_faces: np.ndarray,
    smooth_mask: np.ndarray,
    diagonal: float,
    mean_edge: float,
    original_area: float,
    original_volume: float,
) -> dict[str, float | int]:
    normals, _, _ = face_geometry(vertices, faces, need_centroids=False)
    displacement = np.linalg.norm(vertices - original_vertices, axis=1)
    left, right = adjacent_faces[:, 0], adjacent_faces[:, 1]
    dots = np.einsum("ij,ij->i", normals[left], normals[right])
    normal_delta = np.sqrt(np.maximum(0.0, 2.0 - 2.0 * np.clip(dots, -1.0, 1.0)))
    face_alignment = np.einsum("ij,ij->i", original_normals, normals)
    area = surface_area(vertices, faces)
    volume = abs(signed_volume(vertices, faces))
    return {
        "displacement_mean": float(displacement.mean(dtype=np.float64)),
        "displacement_rms": float(np.sqrt(np.mean(displacement**2, dtype=np.float64))),
        "displacement_p95": float(np.quantile(displacement, 0.95)),
        "displacement_p99": float(np.quantile(displacement, 0.99)),
        "displacement_max": float(displacement.max()),
        "displacement_max_pct_diagonal": float(100.0 * displacement.max() / diagonal),
        "displacement_p95_in_mean_edges": float(np.quantile(displacement, 0.95) / mean_edge),
        "flipped_faces": int(np.count_nonzero(face_alignment <= 0.0)),
        "normal_change_mean_degrees": float(
            np.degrees(np.arccos(np.clip(face_alignment, -1.0, 1.0))).mean(dtype=np.float64)
        ),
        "smooth_region_normal_tv_rms": float(
            np.sqrt(np.mean(np.square(normal_delta[smooth_mask]), dtype=np.float64))
        ),
        "reflection_field_rms": reflection_rms(normals, adjacent_faces, smooth_mask),
        "surface_area": area,
        "surface_area_change_pct": float(100.0 * (area / original_area - 1.0)),
        "absolute_volume": volume,
        "volume_change_pct": float(100.0 * (volume / original_volume - 1.0)),
    }


def main() -> None:
    args = parse_args()
    input_path = Path(args.input).resolve()
    outdir = Path(args.outdir).resolve()
    outdir.mkdir(parents=True, exist_ok=True)
    selected = {config.name: config for config in CONFIGS if config.name in args.variants}

    loaded = np.load(input_path)
    original_vertices = np.asarray(loaded["vertices"], dtype=np.float32)
    faces = np.asarray(loaded["faces"], dtype=np.int32)
    bounds = original_vertices.max(axis=0) - original_vertices.min(axis=0)
    diagonal = float(np.linalg.norm(bounds))
    print(f"Loaded {len(original_vertices):,} vertices / {len(faces):,} triangles")

    started = time.perf_counter()
    topology = build_topology(faces, len(original_vertices))
    unique_edges = topology["unique_edges"]
    adjacent_faces = topology["adjacent_faces"]
    adjacent_edges = topology["adjacent_edges"]
    boundary_vertices = topology["boundary_vertices"]
    edge_length = mean_edge_length(original_vertices, unique_edges)
    print(
        f"Topology: {len(unique_edges):,} edges, {len(adjacent_faces):,} face pairs, "
        f"mean edge {edge_length:.8f}"
    )

    original_face_normals, face_areas, _ = face_geometry(original_vertices, faces)
    original_vertex_normals = vertex_normals(
        original_vertices, faces, face_areas, original_face_normals
    )
    original_area = float(face_areas.sum(dtype=np.float64))
    original_volume = abs(signed_volume(original_vertices, faces))

    max_guide_iterations = max(config.guide_iterations for config in selected.values())
    guidance_levels: dict[int, np.ndarray] = {0: original_face_normals}
    guidance = original_face_normals.copy()
    for iteration in range(1, max_guide_iterations + 1):
        guidance = diffuse_normals(
            guidance,
            face_areas,
            adjacent_faces,
            iterations=1,
            sigma=0.78,
            blend=0.82,
        )
        guidance_levels[iteration] = guidance.copy()

    metric_guidance = guidance_levels[max_guide_iterations]
    metric_left, metric_right = adjacent_faces[:, 0], adjacent_faces[:, 1]
    metric_dots = np.einsum(
        "ij,ij->i", metric_guidance[metric_left], metric_guidance[metric_right]
    )
    metric_distance = np.sqrt(
        np.maximum(0.0, 2.0 - 2.0 * np.clip(metric_dots, -1.0, 1.0))
    )
    # Reflection roughness is compared only where the multiscale guidance says
    # the underlying surface is continuous, excluding intended creases.
    smooth_mask = metric_distance < 0.22

    baseline = geometry_metrics(
        original_vertices,
        original_vertices,
        faces,
        original_face_normals,
        adjacent_faces,
        smooth_mask,
        diagonal,
        edge_length,
        original_area,
        original_volume,
    )
    report: dict[str, object] = {
        "input": str(input_path),
        "vertices": int(len(original_vertices)),
        "triangles": int(len(faces)),
        "bounds_diagonal": diagonal,
        "mean_edge_length": edge_length,
        "smooth_face_pair_fraction": float(np.mean(smooth_mask)),
        "topology": topology["stats"],
        "baseline": baseline,
        "variants": {},
    }

    for name, config in selected.items():
        print(f"[{name}] filtering normals")
        guide = guidance_levels[config.guide_iterations]
        protection, _ = feature_protection(
            guide,
            adjacent_faces,
            adjacent_edges,
            boundary_vertices,
            len(original_vertices),
        )
        target_normals = diffuse_normals(
            original_face_normals,
            face_areas,
            adjacent_faces,
            iterations=config.normal_iterations,
            sigma=config.normal_sigma,
            fixed_guidance=guide,
            blend=0.86,
        )
        print(f"[{name}] projecting original vertices")
        vertices = update_vertices(
            original_vertices,
            faces,
            target_normals,
            original_vertex_normals,
            protection,
            config,
            diagonal,
        )
        vertices, reverted_vertices = revert_unsafe_faces(
            original_vertices, vertices, faces, original_face_normals
        )
        output_path = outdir / f"denoised-{name}.npz"
        np.savez(output_path, vertices=vertices, faces=faces)
        metrics = geometry_metrics(
            original_vertices,
            vertices,
            faces,
            original_face_normals,
            adjacent_faces,
            smooth_mask,
            diagonal,
            edge_length,
            original_area,
            original_volume,
        )
        metrics["normal_tv_improvement_pct"] = float(
            100.0
            * (1.0 - metrics["smooth_region_normal_tv_rms"] / baseline["smooth_region_normal_tv_rms"])
        )
        metrics["reflection_improvement_pct"] = float(
            100.0 * (1.0 - metrics["reflection_field_rms"] / baseline["reflection_field_rms"])
        )
        report["variants"][name] = {
            "config": asdict(config),
            "output": str(output_path),
            "metrics": metrics,
            "reverted_vertices_for_flip_safety": reverted_vertices,
            "protected_vertex_fraction_over_50pct": float(np.mean(protection > 0.5)),
        }
        print(
            f"[{name}] reflection {metrics['reflection_improvement_pct']:.2f}% better; "
            f"p95 displacement {metrics['displacement_p95']:.7f}; "
            f"flips {metrics['flipped_faces']}"
        )
        del target_normals, vertices, protection
        gc.collect()

    report["elapsed_seconds"] = time.perf_counter() - started
    report_path = outdir / "denoise-metrics.json"
    report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(f"Wrote {report_path} in {report['elapsed_seconds']:.1f}s")


if __name__ == "__main__":
    main()