#!/usr/bin/env python3
"""Patch-based robust quadratic surface denoising with a strict trust region.

Each original vertex is projected only along its original normal onto a robust
local quadratic surface fitted to nearby original samples.  Connectivity never
changes.  This removes coherent scan/generative ripples that one-ring filters
cannot see while retaining local curvature better than plane projection.
"""

from __future__ import annotations

import argparse
import json
import time
from pathlib import Path

import numpy as np
from scipy.spatial import cKDTree

import denoise_mesh as dm


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", required=True)
    parser.add_argument("--outdir", required=True)
    return parser.parse_args()


def tangent_frames(normals: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    references = np.zeros_like(normals)
    use_z = np.abs(normals[:, 2]) < 0.88
    references[use_z, 2] = 1.0
    references[~use_z, 0] = 1.0
    tangent_x = dm.normalize_rows(np.cross(references, normals)).astype(np.float32)
    tangent_y = dm.normalize_rows(np.cross(normals, tangent_x)).astype(np.float32)
    return tangent_x, tangent_y


def robust_quadratic_offsets(
    vertices: np.ndarray,
    normals: np.ndarray,
    tree: cKDTree,
    neighbor_count: int,
    *,
    batch_size: int = 1024,
) -> tuple[np.ndarray, dict[str, float]]:
    count = len(vertices)
    offsets = np.zeros(count, dtype=np.float32)
    tangent_x, tangent_y = tangent_frames(normals)
    radius_samples: list[float] = []
    condition_failures = 0

    for start in range(0, count, batch_size):
        end = min(start + batch_size, count)
        centers = vertices[start:end]
        distance, neighbor = tree.query(centers, k=neighbor_count, workers=-1)
        distance = np.asarray(distance, dtype=np.float32)
        neighbor = np.asarray(neighbor, dtype=np.int32)
        delta = vertices[neighbor] - centers[:, None, :]
        nx = tangent_x[start:end]
        ny = tangent_y[start:end]
        nn = normals[start:end]
        patch_radius = np.maximum(distance[:, -1], np.float32(1e-8))
        radius_samples.extend(patch_radius[:: max(1, len(patch_radius) // 16)].tolist())

        x = np.einsum("bki,bi->bk", delta, nx) / patch_radius[:, None]
        y = np.einsum("bki,bi->bk", delta, ny) / patch_radius[:, None]
        z = np.einsum("bki,bi->bk", delta, nn)
        A = np.stack((np.ones_like(x), x, y, x * x, x * y, y * y), axis=2)

        spatial = np.exp(-2.0 * np.square(distance / patch_radius[:, None])).astype(np.float32)
        normal_dot = np.einsum("bki,bi->bk", normals[neighbor], nn)
        normal_distance = np.sqrt(
            np.maximum(0.0, 2.0 - 2.0 * np.clip(normal_dot, -1.0, 1.0))
        )
        normal_weight = np.exp(-0.5 * np.square(normal_distance / 0.38)).astype(np.float32)
        weights = spatial * normal_weight
        weights[:, 0] = np.maximum(weights[:, 0], 1.0)

        def solve(weight: np.ndarray) -> np.ndarray:
            matrix = np.einsum("bki,bk,bkj->bij", A, weight, A, optimize=True)
            rhs = np.einsum("bki,bk,bk->bi", A, weight, z, optimize=True)
            trace = np.trace(matrix, axis1=1, axis2=2)
            regularizer = np.maximum(trace * 1e-7, 1e-10).astype(np.float32)
            matrix[:, range(6), range(6)] += regularizer[:, None]
            # Keep the intercept effectively unregularized so a displaced local
            # surface may move back to the patch rather than toward zero.
            matrix[:, 0, 0] -= regularizer * np.float32(0.999)
            try:
                return np.linalg.solve(matrix, rhs)
            except np.linalg.LinAlgError:
                nonlocal condition_failures
                condition_failures += len(matrix)
                return np.zeros((len(matrix), 6), dtype=np.float32)

        coefficients = solve(weights)
        residual = z - np.einsum("bki,bi->bk", A, coefficients)
        scale = np.maximum(
            np.median(np.abs(residual), axis=1) * np.float32(1.4826),
            patch_radius * np.float32(2e-4),
        )
        robust = np.square(np.clip(1.0 - np.square(residual / (4.685 * scale[:, None])), 0.0, 1.0))
        coefficients = solve(weights * robust.astype(np.float32))
        offsets[start:end] = coefficients[:, 0].astype(np.float32)

    return offsets, {
        "patch_radius_median": float(np.median(radius_samples)),
        "patch_radius_p95": float(np.quantile(radius_samples, 0.95)),
        "condition_failures": condition_failures,
    }


def vertex_area_weights(
    vertex_count: int, faces: np.ndarray, face_areas: np.ndarray
) -> np.ndarray:
    result = np.zeros(vertex_count, dtype=np.float32)
    for corner in range(3):
        result += np.bincount(
            faces[:, corner], weights=face_areas, minlength=vertex_count
        ).astype(np.float32)
    return result / 3.0


def revert_flipped_faces(
    original_vertices: np.ndarray,
    candidate: np.ndarray,
    faces: np.ndarray,
    original_face_normals: np.ndarray,
) -> tuple[np.ndarray, int]:
    vertices = candidate.copy()
    reverted = np.zeros(len(vertices), dtype=bool)
    for _ in range(4):
        normals, areas, _ = dm.face_geometry(vertices, faces, need_centroids=False)
        alignment = np.einsum("ij,ij->i", original_face_normals, normals)
        bad = (alignment <= 0.0) | (areas < 1e-12)
        if not np.any(bad):
            break
        bad_vertices = np.unique(faces[bad].ravel())
        reverted[bad_vertices] = True
        vertices[bad_vertices] = original_vertices[bad_vertices]
    return vertices, int(reverted.sum())


def reflection_metric(
    vertices: np.ndarray,
    faces: np.ndarray,
    adjacent_faces: np.ndarray,
    smooth_mask: np.ndarray,
) -> float:
    normals, _, _ = dm.face_geometry(vertices, faces, need_centroids=False)
    return dm.reflection_rms(normals, adjacent_faces, smooth_mask)


def main() -> None:
    args = parse_args()
    source = Path(args.input).resolve()
    outdir = Path(args.outdir).resolve()
    outdir.mkdir(parents=True, exist_ok=True)
    data = np.load(source)
    original_vertices = np.asarray(data["vertices"], dtype=np.float32)
    faces = np.asarray(data["faces"], dtype=np.int32)
    diagonal = float(
        np.linalg.norm(original_vertices.max(axis=0) - original_vertices.min(axis=0))
    )
    topology = dm.build_topology(faces, len(original_vertices))
    adjacent_faces = topology["adjacent_faces"]
    adjacent_edges = topology["adjacent_edges"]
    boundary_vertices = topology["boundary_vertices"]
    original_face_normals, face_areas, _ = dm.face_geometry(original_vertices, faces)
    original_vertex_normals = dm.vertex_normals(
        original_vertices, faces, face_areas, original_face_normals
    )
    guidance = dm.diffuse_normals(
        original_face_normals, face_areas, adjacent_faces, 12, 0.78, blend=0.82
    )
    feature_strength, guide_distance = dm.feature_protection(
        guidance,
        adjacent_faces,
        adjacent_edges,
        boundary_vertices,
        len(original_vertices),
    )
    smooth_mask = guide_distance < 0.22
    tree = cKDTree(original_vertices)
    vertex_area = vertex_area_weights(len(original_vertices), faces, face_areas)
    original_area = float(face_areas.sum(dtype=np.float64))
    original_volume = abs(dm.signed_volume(original_vertices, faces))
    mean_edge = dm.mean_edge_length(original_vertices, topology["unique_edges"])
    baseline = dm.geometry_metrics(
        original_vertices,
        original_vertices,
        faces,
        original_face_normals,
        adjacent_faces,
        smooth_mask,
        diagonal,
        mean_edge,
        original_area,
        original_volume,
    )

    variants = {
        "patch-small": {"neighbors": 48, "gain": 0.85, "max_pct": 0.0015},
        "patch-balanced": {"neighbors": 96, "gain": 0.95, "max_pct": 0.0030},
        "patch-wide": {"neighbors": 192, "gain": 1.00, "max_pct": 0.0050},
    }
    report: dict[str, object] = {
        "input": str(source),
        "vertices": len(original_vertices),
        "triangles": len(faces),
        "topology": topology["stats"],
        "baseline": baseline,
        "variants": {},
    }
    started = time.perf_counter()
    for name, config in variants.items():
        print(f"[{name}] robust quadratic patches ({config['neighbors']} neighbors)")
        offsets, patch_stats = robust_quadratic_offsets(
            original_vertices, original_vertex_normals, tree, config["neighbors"]
        )
        offsets *= np.float32(config["gain"]) * (1.0 - 0.97 * feature_strength)
        # First-order volume preservation for a closed oriented mesh.
        mean_offset = float(np.sum(offsets * vertex_area) / np.maximum(vertex_area.sum(), 1e-20))
        offsets -= np.float32(mean_offset) * (1.0 - feature_strength)
        trust = np.float32(diagonal * config["max_pct"])
        offsets = np.clip(offsets, -trust, trust)
        candidate = original_vertices + original_vertex_normals * offsets[:, None]
        candidate, reverted_vertices = revert_flipped_faces(
            original_vertices, candidate, faces, original_face_normals
        )
        output = outdir / f"{name}.npz"
        np.savez(output, vertices=candidate.astype(np.float32), faces=faces)
        metrics = dm.geometry_metrics(
            original_vertices,
            candidate,
            faces,
            original_face_normals,
            adjacent_faces,
            smooth_mask,
            diagonal,
            mean_edge,
            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": config,
            "output": str(output),
            "patch": patch_stats,
            "raw_offset_mean": float(offsets.mean()),
            "raw_offset_rms": float(np.sqrt(np.mean(offsets * offsets))),
            "raw_offset_p95_abs": float(np.quantile(np.abs(offsets), 0.95)),
            "reverted_vertices_for_flip_safety": reverted_vertices,
            "metrics": metrics,
        }
        print(
            f"[{name}] reflection {metrics['reflection_improvement_pct']:.2f}% better, "
            f"p95 {metrics['displacement_p95']:.7f}, flips {metrics['flipped_faces']}"
        )
    report["elapsed_seconds"] = time.perf_counter() - started
    (outdir / "patch-metrics.json").write_text(
        json.dumps(report, indent=2), encoding="utf-8"
    )
    print(f"Done in {report['elapsed_seconds']:.1f}s")


if __name__ == "__main__":
    main()