#!/usr/bin/env python3
"""Create conservative custom vertex normals without moving mesh geometry."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

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 target_vertex_normals(
    faces: np.ndarray,
    face_areas: np.ndarray,
    filtered_face_normals: np.ndarray,
    vertex_count: int,
) -> np.ndarray:
    weighted = filtered_face_normals * face_areas[:, None]
    normals = np.zeros((vertex_count, 3), dtype=np.float32)
    for corner in range(3):
        indices = faces[:, corner]
        for axis in range(3):
            normals[:, axis] += np.bincount(
                indices,
                weights=weighted[:, axis],
                minlength=vertex_count,
            ).astype(np.float32)
    return dm.normalize_rows(normals).astype(np.float32)


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


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)
    vertices = np.asarray(data["vertices"], dtype=np.float32)
    faces = np.asarray(data["faces"], dtype=np.int32)
    topology = dm.build_topology(faces, len(vertices))
    adjacent_faces = topology["adjacent_faces"]
    adjacent_edges = topology["adjacent_edges"]
    boundary_vertices = topology["boundary_vertices"]
    edges = topology["unique_edges"]
    face_normals, face_areas, _ = dm.face_geometry(vertices, faces)
    base_vertex_normals = dm.vertex_normals(vertices, faces, face_areas, face_normals)

    report: dict[str, object] = {
        "input": str(source),
        "vertices": len(vertices),
        "triangles": len(faces),
        "baseline_reflection_vertex_rms": reflection_vertex_rms(base_vertex_normals, edges),
        "variants": {},
    }
    variants = {
        "normal-soft": (4, 10, 0.34, 0.28),
        "normal-balanced": (8, 20, 0.44, 0.52),
        "normal-gloss": (14, 36, 0.56, 0.72),
    }
    for name, (guide_iterations, iterations, sigma, blend) in variants.items():
        guidance = dm.diffuse_normals(
            face_normals, face_areas, adjacent_faces, guide_iterations, 0.78, blend=0.82
        )
        protection, _ = dm.feature_protection(
            guidance, adjacent_faces, adjacent_edges, boundary_vertices, len(vertices)
        )
        filtered_faces = dm.diffuse_normals(
            face_normals,
            face_areas,
            adjacent_faces,
            iterations,
            sigma,
            fixed_guidance=guidance,
            blend=0.88,
        )
        target = target_vertex_normals(faces, face_areas, filtered_faces, len(vertices))
        blend_per_vertex = np.float32(blend) * (1.0 - 0.94 * protection)
        custom = dm.normalize_rows(
            (1.0 - blend_per_vertex[:, None]) * base_vertex_normals
            + blend_per_vertex[:, None] * target
        ).astype(np.float32)
        output = outdir / f"{name}.npz"
        np.savez(output, vertices=vertices, faces=faces, normals=custom)
        angular_change = np.degrees(
            np.arccos(np.clip(np.einsum("ij,ij->i", base_vertex_normals, custom), -1.0, 1.0))
        )
        reflection = reflection_vertex_rms(custom, edges)
        report["variants"][name] = {
            "output": str(output),
            "guide_iterations": guide_iterations,
            "normal_iterations": iterations,
            "normal_sigma": sigma,
            "blend": blend,
            "normal_change_mean_degrees": float(angular_change.mean()),
            "normal_change_p95_degrees": float(np.quantile(angular_change, 0.95)),
            "reflection_vertex_rms": reflection,
            "reflection_improvement_pct": float(
                100.0 * (1.0 - reflection / report["baseline_reflection_vertex_rms"])
            ),
        }
        print(name, json.dumps(report["variants"][name], indent=2))
    (outdir / "shading-normal-metrics.json").write_text(
        json.dumps(report, indent=2), encoding="utf-8"
    )


if __name__ == "__main__":
    main()