"""Noisy versions of the report's meshes, topology kept: ../meshes/<name>_noise<pct>.ply for sigma = 0.2, 0.6, 1 % of the
bounding-box diagonal, and ../report/noisy.html (a validation table and renders).

Each mesh is first subdivided to edges of at most MAX_EDGE of the diagonal (same shape: a coarse mesh would otherwise
only move its few big faces). Every vertex x then moves by u(x), a smooth random 3D vector field: per component, white
noise on a grid Gaussian-filtered to a correlation length of CORR sigma, read at the vertices (trilinear), scaled to a
standard deviation of sigma (so the displacement along any direction, the normal included, has std sigma). The
correlation length is raised until the field's largest gradient (Frobenius norm over the grid, an upper bound of its
Lipschitz constant) is below LIP: then x -> x + u(x) is one-to-one in all of space, so the surface cannot cross itself
or turn a face over, and a thin wall moves as a whole. Faces are unchanged (same connectivity, watertight). Checked:
no face turned over, and the exact sections at the 128 layer centres along up are simple rings.

Tried first and dropped: per-vertex noise along the normal folds faces at this size, and noise along the normal (even
smooth) makes surfaces cross at sharp edges: a wall moving in and the floor next to it moving up meet.

    env -u PYTHONPATH PYTHONWARNINGS=ignore .venv/bin/python ~/code/data/meshes/fcp_2710_tests/scripts/make_noisy.py
"""
import json
import sys
import time
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
import trimesh
from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from scipy import ndimage
from shapely.geometry import LinearRing

sys.path.insert(0, str(Path(__file__).parent))
from parts_report import CSS, OBJ, REPORT, ROOT, UP

SIGMAS = (0.002, 0.006, 0.010)  # of the diagonal
MAX_EDGE = 0.005  # of the diagonal
CORR = 5.0  # starting correlation length of the field, in sigmas
LIP = 0.8  # largest gradient of the field allowed (below 1: x -> x + u(x) is one-to-one)
OUT = ROOT / "meshes"


def label(s):
    return f"{100 * s:g}".replace(".", "p")  # 0.002 -> "0p2"


def face_normals(V, F):
    n = np.cross(V[F[:, 1]] - V[F[:, 0]], V[F[:, 2]] - V[F[:, 0]])
    return n / np.maximum(np.linalg.norm(n, axis=1, keepdims=True), 1e-300)


def sections_ok(mesh, up, layers=128):
    """(rings, rings that self-cross) over the exact sections at the layer centres along up."""
    up = np.asarray(up, float)
    h = mesh.vertices @ up
    lo, hi = h.min(), h.max()
    zs = lo + (np.arange(layers) + 0.5) * (hi - lo) / layers
    sec = mesh.section_multiplane(plane_origin=lo * up, plane_normal=up, heights=zs - lo)
    n = bad = 0
    for s in sec:
        if s is None:
            continue
        for L in s.discrete:
            if len(L) < 4:
                continue
            n += 1
            bad += not LinearRing(L).is_simple
    return n, bad


def vector_field(V, std, length, rng):
    """u at the points V: a smooth random 3D field, each component of standard deviation std, Gaussian correlation
    `length`; and its largest gradient (Frobenius norm over the grid)."""
    h = length / 3  # voxel
    lo = V.min(0) - 3 * length
    shape = tuple(np.ceil((V.max(0) + 3 * length - lo) / h).astype(int) + 1)
    G = [ndimage.gaussian_filter(rng.standard_normal(shape).astype(np.float32), sigma=length / h, mode="wrap")
         for _ in range(3)]
    G = [g * (std / g.std()) for g in G]
    lip = float(np.sqrt(sum(gr ** 2 for g in G for gr in np.gradient(g, h))).max())
    u = np.stack([ndimage.map_coordinates(g, ((V - lo) / h).T, order=1) for g in G], 1)
    return u, lip


def noisy(mesh, sigma, diag, rng):
    """(noisy mesh, displacements u, correlation length used / diag, the field's largest gradient, faces turned over)."""
    length = CORR * sigma * diag
    while True:
        u, lip = vector_field(mesh.vertices, sigma * diag, length, rng)
        if lip < LIP:
            break
        length *= 1.25
    from scipy.sparse import coo_matrix
    from scipy.sparse.csgraph import connected_components

    F, n0, E = mesh.faces, mesh.face_normals, mesh.edges_unique
    # a near-collinear sliver of the source mesh can still turn over (the field's curvature across it exceeds its
    # height): the vertices of turned faces, grouped into connected patches, move rigidly by their patch's mean u (a
    # translation turns no face inside the patch); a face at a patch's edge that turns joins it on the next round
    u0, S = u.copy(), np.zeros(len(u), bool)
    while (t := (face_normals(mesh.vertices + u, F) * n0).sum(1) <= 0).any():
        S[F[t].ravel()] = True
        e = E[S[E[:, 0]] & S[E[:, 1]]]
        _, lab = connected_components(coo_matrix((np.ones(len(e)), (e[:, 0], e[:, 1])), shape=(len(u),) * 2), directed=False)
        u = u0.copy()
        for c in np.unique(lab[S]):
            vs = np.nonzero(S & (lab == c))[0]
            u[vs] = u0[vs].mean(0)
    V = mesh.vertices + u
    turned = int(((face_normals(V, F) * n0).sum(1) <= 0).sum())
    return trimesh.Trimesh(V, F, process=False), u, length / diag, lip, turned


def render(meshes, titles, path):
    fig = plt.figure(figsize=(4 * len(meshes), 4.4), dpi=100)
    light = np.array([0.45, -0.6, 0.65]); light /= np.linalg.norm(light)
    allv = meshes[0].vertices
    lo, hi = allv.min(0), allv.max(0)
    m, r = (lo + hi) / 2, (hi - lo).max() / 2
    for k, (mesh, t) in enumerate(zip(meshes, titles)):
        ax = fig.add_subplot(1, len(meshes), k + 1, projection="3d", proj_type="ortho")
        fn = mesh.face_normals
        shade = 0.25 + 0.75 * np.abs(fn @ light)
        ax.add_collection3d(Poly3DCollection(mesh.triangles, facecolors=np.c_[shade, shade, shade] * [0.85, 0.82, 0.75],
                                             edgecolors="none", linewidths=0))
        ax.set_xlim(m[0] - r, m[0] + r); ax.set_ylim(m[1] - r, m[1] + r); ax.set_zlim(m[2] - r, m[2] + r)
        ax.set_box_aspect((1, 1, 1)); ax.set_axis_off(); ax.view_init(15, -60)
        ax.set_title(t, fontsize=9)
    fig.tight_layout()
    fig.savefig(path)
    plt.close(fig)


def main():
    OUT.mkdir(exist_ok=True)
    rows = []
    for name, p in OBJ.items():
        t0 = time.perf_counter()
        g = trimesh.load(p, force="mesh")
        diag = float(np.linalg.norm(g.extents))
        v, f = trimesh.remesh.subdivide_to_size(g.vertices, g.faces, MAX_EDGE * diag, max_iter=20)
        base = trimesh.Trimesh(v, f, process=True)
        assert base.is_watertight and base.is_winding_consistent, name
        n = base.vertex_normals
        up = UP.get(name, (0, 0, 1))
        r0, b0 = sections_ok(base, up)
        meshes, titles = [base], [f"clean (subdivided, {len(base.faces):,} faces)"]
        for i, s in enumerate(SIGMAS):
            m, u, corr, lip, turned = noisy(base, s, diag, np.random.default_rng(1000 + i))
            m.export(OUT / f"{name}_noise{label(s)}.ply")
            rings, bad = sections_ok(m, up)
            d = (u * n).sum(1)  # the displacement along the normal
            row = dict(name=name, sigma=s, faces=len(m.faces), watertight=bool(m.is_watertight),
                       volume_ratio=float(m.volume / base.volume), std=float(d.std() / diag), corr=corr, lip=lip,
                       turned=turned, rings=rings, rings_clean=r0, bad_rings=bad, bad_clean=b0)
            rows.append(row)
            meshes.append(m)
            titles.append(f"σ = {100 * s:g} % diag (along the normal {100 * row['std']:.2f} %)")
            print(f"{name:16s} sigma {100 * s:4.1f}%  normal std {100 * row['std']:.3f}%  corr {100 * corr:.1f}%  "
                  f"lip {lip:.2f}  turned {turned}  self-crossing rings {bad}/{rings} (clean {b0}/{r0})  "
                  f"vol {row['volume_ratio']:.4f}", flush=True)
        render(meshes, titles, REPORT / f"img/noisy_{name}.png")
        print(f"  {name}: {time.perf_counter() - t0:.1f} s", flush=True)
    js = OUT / "noisy_meshes.json"  # merged with the other objects' rows (a run may cover a few objects: FCP_OBJECTS)
    done = {r["name"] for r in rows}
    old = [r for r in json.loads(js.read_text()) if r["name"] not in done and r["name"] != "loaf_pan"] if js.exists() else []
    rows = old + rows
    js.write_text(json.dumps(rows, indent=1))
    write_html(rows)


def write_html(rows):
    tr = "".join(
        f"<tr><td>{r['name']}</td><td>{100 * r['sigma']:g} %</td><td>{100 * r['std']:.3f} %</td><td>{100 * r['corr']:.1f} %</td>"
        f"<td>{r['lip']:.2f}</td><td>{r['faces']:,}</td><td>{'yes' if r['watertight'] else '<b>no</b>'}</td><td>{r['turned']}</td>"
        f"<td class='{'bad' if r['bad_rings'] > r['bad_clean'] else ''}'>{r['bad_rings']} / {r['rings']}</td>"
        f"<td>{r['bad_clean']} / {r['rings_clean']}</td><td>{r['volume_ratio']:.4f}</td></tr>" for r in rows)
    imgs = "".join(f"<h2 id='{n}'>{n}</h2><img src='img/noisy_{n}.png' loading='lazy'>" for n in dict.fromkeys(r["name"] for r in rows))
    page = f"""<!doctype html><html lang="en"><head><meta charset="utf-8">
<meta name="viewport" content="width=device-width,initial-scale=1"><title>Noisy meshes</title><style>{CSS}</style></head><body>
<p><a href="index.html">← report</a></p>
<h1>Noisy meshes</h1>
<p class="mut">Files: <code>meshes/&lt;name&gt;_noise0p2.ply</code>, <code>_noise0p6</code>, <code>_noise1.ply</code>; script <code>scripts/make_noisy.py</code>.
Each mesh is subdivided to edges ≤ {100 * MAX_EDGE:g} % of the diagonal (same shape), then every vertex x moves by u(x), a smooth random
3D vector field (seeded): each component Gaussian with std σ = 0.2 / 0.6 / 1 % of the diagonal and a Gaussian correlation length of at
least {CORR:g} σ, raised until the field's largest gradient is below {LIP:g}. Then x → x + u(x) is one-to-one, so the surface cannot
cross itself or turn a face over, whatever the wall thickness or the edges. Faces are unchanged (same connectivity, watertight).
"along the normal": the achieved displacement along the vertex normals, over the diagonal. Self-crossing rings: the exact sections
at the 128 layer centres along up that are not simple (red: more than the clean mesh). Generated {time.strftime('%Y-%m-%d %H:%M')}.</p>
<table><tr><th>object</th><th>σ</th><th>along the normal</th><th>correlation length</th><th>max gradient</th><th>faces</th>
<th>watertight</th><th>faces turned over</th><th>self-crossing rings</th><th>clean</th><th>volume / clean</th></tr>{tr}</table>
{imgs}
</body></html>"""
    (REPORT / "noisy.html").write_text(page)


if __name__ == "__main__":
    main()