"""Volume: exact (the mesh), FCP from the mesh, FCP from a 20k point cloud, screened Poisson from the same cloud (true
normals, and normals estimated from the points). 128 layers, 2048 atoms. For the report section "volume".

    cd ~/code/data/meshes/fcp_3010_hull
    F=~/code/fcp_devel/.venv/bin/python; O=~/code/data/meshes/fcp_2910_collision/.venv/bin/python
    env -u PYTHONPATH $F scripts/volume_pc.py points      # out/pc/<obj>.npz: 20k area-uniform points, seed 3, true normals
    env -u PYTHONPATH $F scripts/volume_study.py fcp      # out/vs_fcp.csv, out/vs_render/<obj>_fcp_*.npz
    env -u PYTHONPATH OMP_NUM_THREADS=1 $O scripts/volume_study.py poisson   # out/vs_poisson.csv, out/vs_render/<obj>_poisson_*.npz
    env -u PYTHONPATH $O scripts/volume_study.py render   # sections/hull/volume_fig_renders.png (Open3D decimates for drawing)
    env -u PYTHONPATH $F scripts/volume_study.py section  # sections/hull/{volume_data.csv, volume_objects.csv, volume_section.html}

Times: median of 5 (encode, Poisson) or of 2000 (volumes) after a warm-up. FCP's CPU encode uses numba's threads (12);
Poisson runs on one thread (Open3D's threaded Poisson was 1.5-4x slower on this CPU).
"""
import json
import sys
import time
from pathlib import Path

import numpy as np

STEP = sys.argv[1]
sys.argv = [sys.argv[0], "x"]
sys.path.insert(0, str(Path(__file__).parent))
from hull import OBJECTS  # noqa: E402

ROOT = Path(__file__).resolve().parent.parent
OUT, PC, RD = ROOT / "out", ROOT / "out/pc", ROOT / "out/vs_render"
SECTION = Path.home() / "code/data/meshes/fcp_report/sections/hull"  # the report session merges; files named volume_*
import os  # noqa: E402
LAYERS = int(os.environ.get("FCP_LAYERS", 128))
ATOMS = int(os.environ.get("FCP_ATOMS", 2048))
# the base run (128 layers, 2048 atoms): out/vs_fcp.csv, columns fcp_*; the others: out/vs_fcp<TAG>.csv, columns <prefix>_*
VARIANTS = [("_6144", "fcp6144", "128 layers, 6144 atoms"), ("_L256_6144", "fcpL256", "256 layers, 6144 atoms")]
TAG = "" if (LAYERS, ATOMS) == (128, 2048) else (f"_{ATOMS}" if LAYERS == 128 else f"_L{LAYERS}_{ATOMS}")
RENDER = ["mug", "bowl_dog", "shoe", "dino", "horse", "knot"]


def med(f, n=5):
    f()
    t = []
    for _ in range(n):
        t0 = time.perf_counter()
        f()
        t.append(time.perf_counter() - t0)
    return 1e3 * float(np.median(t))


def mesh_volume(V, F):
    a, b, c = V[F[:, 0]], V[F[:, 1]], V[F[:, 2]]
    return float(np.einsum("ij,ij->", a, np.cross(b, c)) / 6)


def fcp_step():
    import pandas as pd
    import trimesh

    import fcp
    from volume_atoms import volume_atoms_fast
    from volume_sota import vol_nb  # numba tetrahedron sum

    RD.mkdir(parents=True, exist_ok=True)
    rows = []
    for name, path in OBJECTS.items():
        m = trimesh.load(path, force="mesh")
        V0, F0 = np.ascontiguousarray(m.vertices, np.float64), np.ascontiguousarray(m.faces, np.int64)
        exact = vol_nb(V0, F0)
        P = np.asarray(np.load(PC / f"{name}.npz")["P"], np.float32)
        r = dict(obj=name, faces=len(F0), exact=exact, exact_us=1e3 * med(lambda: vol_nb(V0, F0), 2000))
        for src, geo in (("mesh", m), ("points", P)):
            d = fcp.encode(geo, atoms=ATOMS, layers=LAYERS, device="cpu")
            r[f"fcp_{src}_vol"] = volume_atoms_fast(d)
            r[f"fcp_{src}_vol_us"] = 1e3 * med(lambda: volume_atoms_fast(d), 2000)
            r[f"fcp_{src}_enc_cpu_ms"] = med(lambda: fcp.encode(geo, atoms=ATOMS, layers=LAYERS, device="cpu"))
            r[f"fcp_{src}_enc_cuda_ms"] = med(lambda: fcp.encode(geo, atoms=ATOMS, layers=LAYERS, device="cuda"))
            r[f"fcp_{src}_bytes"] = d.nbytes
            dc = fcp.encode(geo, atoms=ATOMS, layers=LAYERS, device="cuda")
            r[f"fcp_{src}_vol_cuda"] = volume_atoms_fast(dc)  # the CUDA encode's descriptor (round-off apart)
            if name in RENDER and not TAG:
                dm = fcp.decode(d, device="cpu")
                np.savez(RD / f"{name}_fcp_{src}.npz", V=dm.vertices, F=dm.faces)
        if name in RENDER and not TAG:
            np.savez(RD / f"{name}_mesh.npz", V=V0, F=F0)
        rows.append(r)
        print({k: (round(v, 4) if isinstance(v, float) else v) for k, v in r.items()}, flush=True)
    pd.DataFrame(rows).to_csv(OUT / f"vs_fcp{TAG}.csv", index=False)


def poisson_step():
    import open3d as o3d
    import pandas as pd

    rows = []
    for name in OBJECTS:
        z = np.load(PC / f"{name}.npz")
        r = dict(obj=name)
        for kind in ("true", "est"):
            def build():
                pc = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(z["P"]))
                if kind == "true":
                    pc.normals = o3d.utility.Vector3dVector(z["N"])
                else:  # what a scan gives: normals from 30 neighbours, oriented by the tangent-plane graph
                    pc.estimate_normals(o3d.geometry.KDTreeSearchParamKNN(30))
                    pc.orient_normals_consistent_tangent_plane(30)
                return pc

            pc = build()
            m, _ = o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(pc, depth=8, n_threads=1)
            V, F = np.asarray(m.vertices), np.asarray(m.triangles)
            r[f"poisson_{kind}_vol"] = abs(mesh_volume(V, F))  # a globally flipped orientation counts as right
            r[f"poisson_{kind}_watertight"] = bool(m.is_watertight())
            r[f"poisson_{kind}_ms"] = med(lambda: o3d.geometry.TriangleMesh.create_from_point_cloud_poisson(
                build() if kind == "est" else pc, depth=8, n_threads=1), 3)
            r[f"poisson_{kind}_vol_us"] = 1e3 * med(lambda: mesh_volume(V, F), 200)
            if name in RENDER:
                np.savez(RD / f"{name}_poisson_{kind}.npz", V=V, F=F)
        rows.append(r)
        print(r, flush=True)
    pd.DataFrame(rows).to_csv(OUT / "vs_poisson.csv", index=False)


COLS = [("mesh", "original mesh"), ("fcp_mesh", "FCP from the mesh"), ("fcp_points", "FCP from 20k points"),
        ("poisson_true", "Poisson, true normals"), ("poisson_est", "Poisson, estimated normals")]


def render_step():
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    import open3d as o3d
    import pandas as pd
    from mpl_toolkits.mplot3d.art3d import Poly3DCollection

    t = table()
    light = np.array([0.3, -0.5, 0.8]) / np.linalg.norm([0.3, -0.5, 0.8])
    tint = {"mesh": (0.86, 0.86, 0.88), "fcp_mesh": (0.55, 0.75, 0.95), "fcp_points": (0.45, 0.65, 0.95),
            "poisson_true": (0.95, 0.75, 0.5), "poisson_est": (0.95, 0.6, 0.45)}
    fig = plt.figure(figsize=(3.6 * len(COLS), 3.5 * len(RENDER)))
    for i, name in enumerate(RENDER):
        z0 = np.load(RD / f"{name}_mesh.npz")
        lo, hi = z0["V"].min(0), z0["V"].max(0)
        c, r = (lo + hi) / 2, (hi - lo).max() / 2
        row = t[t.obj == name].iloc[0]
        for j, (key, label) in enumerate(COLS):
            z = np.load(RD / f"{name}_{key}.npz")
            o = o3d.geometry.TriangleMesh(o3d.utility.Vector3dVector(z["V"]), o3d.utility.Vector3iVector(z["F"]))
            if len(z["F"]) > 25000:
                o = o.simplify_quadric_decimation(25000)
            o.compute_triangle_normals()
            V, F, N = np.asarray(o.vertices), np.asarray(o.triangles), np.asarray(o.triangle_normals)
            sh = 0.35 + 0.6 * np.abs(N @ light)
            ax = fig.add_subplot(len(RENDER), len(COLS), i * len(COLS) + j + 1, projection="3d")
            ax.add_collection3d(Poly3DCollection(V[F], facecolors=np.c_[sh[:, None] * np.array(tint[key])], edgecolors="none"))
            ax.set_xlim(c[0] - r, c[0] + r)
            ax.set_ylim(c[1] - r, c[1] + r)
            ax.set_zlim(c[2] - r, c[2] + r)
            ax.set_box_aspect((1, 1, 1))
            ax.view_init(22, -60)
            ax.set_axis_off()
            txt = "exact" if key == "mesh" else f"{row[key + '_err']:+.2f} %"
            if i == 0:
                ax.set_title(label, fontsize=13)
            ax.text2D(0.5, 0.02, f"{name.replace('_', ' ')}: {txt}" if j == 0 else txt, transform=ax.transAxes,
                      ha="center", fontsize=12)
    fig.subplots_adjust(left=0, right=1, top=0.97, bottom=0, wspace=0, hspace=0)
    SECTION.mkdir(parents=True, exist_ok=True)
    fig.savefig(SECTION / "volume_fig_renders.png", dpi=90, facecolor="white")


def table():
    import pandas as pd

    t = pd.read_csv(OUT / "vs_fcp.csv").merge(pd.read_csv(OUT / "vs_poisson.csv"), on="obj")
    keys = ["fcp_mesh", "fcp_points", "poisson_true", "poisson_est"]
    for tag, pre, _ in VARIANTS:
        if (OUT / f"vs_fcp{tag}.csv").exists():
            b = pd.read_csv(OUT / f"vs_fcp{tag}.csv")
            b = b[["obj"] + [c for c in b if c.startswith("fcp_")]].rename(columns=lambda c, pre=pre: c.replace("fcp_", pre + "_", 1))
            t = t.merge(b, on="obj")
            keys += [pre + "_mesh", pre + "_points"]
    for k in keys:
        t[k + "_err"] = 100 * (t[k + "_vol"] / t.exact - 1)
    return t


def section_step():
    t = table()
    t.to_csv(SECTION / "volume_objects.csv", index=False)
    f = lambda v, n=2: "&ndash;" if v != v else f"{v:.{n}f}"  # noqa: E731
    am = lambda k: (t[k + "_err"].abs().median(), t[k + "_err"].abs().max())  # noqa: E731
    methods = [  # label, error key, build time (ms) and its note, volume time (us)
        ("exact: tetrahedron sum over the source mesh", None, None, "", t.exact_us.median()),
        ("FCP from the mesh (encode 2048 atoms, 128 layers)", "fcp_mesh",
         t.fcp_mesh_enc_cpu_ms.median(), f"{f(t.fcp_mesh_enc_cuda_ms.median(), 0)} ms on CUDA", t.fcp_mesh_vol_us.median()),
        ("FCP from 20k points (same encode, no normals)", "fcp_points",
         t.fcp_points_enc_cpu_ms.median(), f"{f(t.fcp_points_enc_cuda_ms.median(), 0)} ms on CUDA", t.fcp_points_vol_us.median()),
    ] + [row for tag, pre, lab in VARIANTS if pre + "_mesh_err" in t for row in (
        (f"FCP from the mesh ({lab})", pre + "_mesh", t[pre + "_mesh_enc_cpu_ms"].median(),
         f"{f(t[pre + '_mesh_enc_cuda_ms'].median(), 0)} ms on CUDA", t[pre + "_mesh_vol_us"].median()),
        (f"FCP from 20k points ({lab})", pre + "_points", t[pre + "_points_enc_cpu_ms"].median(),
         f"{f(t[pre + '_points_enc_cuda_ms'].median(), 0)} ms on CUDA", t[pre + "_points_vol_us"].median()))] + [
        ("screened Poisson (depth 8) from 20k points, true normals", "poisson_true", t.poisson_true_ms.median(), "1 thread",
         1e3 * t.poisson_true_vol_us.median() / 1e3),
        ("screened Poisson from 20k points, normals estimated (30 neighbours, tangent-plane orientation)", "poisson_est",
         t.poisson_est_ms.median(), "1 thread, normals included", t.poisson_est_vol_us.median()),
    ]
    rows, data = [], []
    for label, k, build, note, vus in methods:
        e = am(k) if k else (0.0, 0.0)
        within = int((t[k + "_err"].abs() <= 1).sum()) if k else len(t)
        rows.append(f"<tr><td style='text-align:left'>{label}</td><td>{'exact' if not k else f(e[0])}</td><td>{'exact' if not k else f(e[1])}</td>"
                    f"<td>{within} / {len(t)}</td><td>{'&ndash;' if build is None else f(build, 0)}{'' if not note else f' ({note})'}</td>"
                    f"<td>{f(vus, 1)}</td></tr>")
        data.append(dict(method=k or "exact", err_med_pct=e[0], err_max_pct=e[1], within_1pct=within, build_ms_median=build,
                         volume_us_median=vus))
    import pandas as pd

    pd.DataFrame(data).to_csv(SECTION / "volume_data.csv", index=False)
    orows = "".join(
        f"<tr><td style='text-align:left'>{r.obj.replace('_', ' ')}</td><td>{r.faces}</td>"
        f"<td>{f(r.fcp_mesh_err)}</td><td>{f(r.fcp_points_err)}</td>"
        + "".join(f"<td>{f(getattr(r, pre + '_mesh_err'))}</td><td>{f(getattr(r, pre + '_points_err'))}</td>" for _, pre, _ in VARIANTS) +
        f"<td>{f(r.poisson_true_err)}</td><td>{f(r.poisson_est_err)}</td>"
        f"<td>{f(r.exact_us, 0)}</td><td>{f(r.fcp_points_vol_us, 0)}</td><td>{f(r.fcp_mesh_enc_cpu_ms, 0)}</td><td>{f(r.fcp_points_enc_cpu_ms, 0)}</td>"
        f"<td>{f(r.fcp_points_enc_cuda_ms, 0)}</td><td>{f(r.poisson_true_ms, 0)}</td><td>{f(r.poisson_est_ms, 0)}</td></tr>"
        for r in t.itertuples())
    fm, fp, pt, pe = am("fcp_mesh"), am("fcp_points"), am("poisson_true"), am("poisson_est")
    bad = t[t.poisson_est_err.abs() > 5].sort_values("poisson_est_err", ascending=False)
    badtxt = ", ".join(f"{o.replace('_', ' ')} {e:+.0f} %" for o, e in zip(bad.obj, bad.poisson_est_err))
    html = f"""<h2>Volume straight from the descriptor</h2>
<p>Package v5.31 (commit db3e6ceb; the working tree only raises <code>MAX_ATOMS</code>), defaults
(<code>parts="tracked-fit"</code>, <code>interp="poly"</code>, the gwn point slicer), 128 layers, 2048 atoms; the 12 objects
of the collision section. The exact volume is the source mesh's (tetrahedra over its faces). From each mesh, 20k
area-uniform surface points (seed 3) with the true normals; every point method gets the same points.
<b>FCP's volume is one contraction of the stored atoms</b>: a ring's area is &pi;&nbsp;&Sigma;<sub>k</sub>&nbsp;k&nbsp;(a<sub>x,k</sub>b<sub>y,k</sub>&nbsp;&minus;&nbsp;b<sub>x,k</sub>a<sub>y,k</sub>)
in its Fourier coefficients, and the height atoms are orthonormal on a part's layer centres, so the sum of the ring areas
over the layers is the same expression in the atom vectors: V&nbsp;=&nbsp;dz&nbsp;&middot;&nbsp;&prod;scale&nbsp;&middot;&nbsp;&pi;&nbsp;&Sigma;<sub>parts</sub>&nbsp;sign&nbsp;&middot;&nbsp;|&Sigma;<sub>k</sub>&nbsp;k&nbsp;(&lang;a<sub>x,k</sub>,&nbsp;b<sub>y,k</sub>&rang;&nbsp;&minus;&nbsp;&lang;b<sub>x,k</sub>,&nbsp;a<sub>y,k</sub>&rang;)|
(the midpoint rule over the layers; no ring is evaluated, no grid, no mesh; one numba loop). Standard route from points:
screened Poisson reconstruction (Open3D 0.20, depth 8), then its mesh's volume, with the true normals (the best case) and
with normals estimated from the points (what a scan gives). CPU: Intel i5-10400F; FCP's encode on numba's 12 threads and on a
GTX 1650; Poisson on one thread (its threaded run was slower here). Times are medians over the objects.</p>
<div class="tablewrap"><table>
<thead><tr><th style='text-align:left'>method</th><th>volume error, median |%|</th><th>max |%|</th><th>objects within 1 %</th>
<th>building the representation, ms</th><th>the volume, &micro;s</th></tr></thead>
<tbody>{''.join(rows)}</tbody></table></div>
<p><b>Accuracy:</b> FCP is within {f(fm[1])} % of the exact volume from the mesh (median {f(fm[0])} %) and within {f(fp[1])} % from the
points (median {f(fp[0])} %), without normals; Poisson with the true normals is as close (median {f(pt[0])} %, max {f(pt[1])} %). With
estimated normals Poisson is off by more than 5 % on {len(bad)} of {len(t)} objects ({badtxt}): on the thin-walled ones the
orientation flips across a wall and the surface closes over the opening (renders). <b>More atoms, more layers:</b> 6144 atoms at 128 layers and at 256 layers move the error little:
from the mesh the median falls {f(fm[0])} &rarr; {f(am("fcp6144_mesh")[0])} &rarr; {f(am("fcpL256_mesh")[0])} % but the worst stays
{f(am("fcpL256_mesh")[1])} %; from the points it is {f(fp[0])} / {f(am("fcp6144_points")[0])} / {f(am("fcpL256_points")[0])} %. The per-object errors keep their sign
and size across the settings (mugs and bowls 0.2&ndash;0.7 % low, the spring from points 1.8 % high): a bias of the
encoding and of the point slicer, not a shortage of atoms. The 6144-atom runs had other jobs on the machine (load average
about 7, against about 0 for the 2048 run): their times in the first table are upper bounds. <b>Time:</b> once encoded, FCP's volume takes {f(t.fcp_points_vol_us.min(), 0)}&ndash;{f(t.fcp_points_vol_us.max(), 0)} &micro;s,
about the cost of the tetrahedron sum over the source mesh itself ({f(t.exact_us.min(), 0)}&ndash;{f(t.exact_us.max(), 0)} &micro;s, numba):
faster on {int((t.fcp_points_vol_us < t.exact_us).sum())} of {len(t)} objects, the gain growing with the mesh (spring, 77k faces: {f(t.exact_us.max() / t[t.exact_us == t.exact_us.max()].fcp_points_vol_us.iloc[0], 1)}&times;);
the volume of a Poisson mesh takes {f(t.poisson_true_vol_us.min() / 1e3, 1)}&ndash;{f(t.poisson_true_vol_us.max() / 1e3, 1)} ms (NumPy). The cost from points is the encode:
{f(t.fcp_points_enc_cpu_ms.median() / 1e3, 1)} s on the CPU and {f(t.fcp_points_enc_cuda_ms.median() / 1e3, 1)} s on CUDA (median), against {f(t.poisson_true_ms.median() / 1e3, 1)} s for Poisson
on one thread; the descriptor, once made, also answers the other queries of this report. Caveats: clean, complete,
uniformly sampled points; Poisson's meshes are not watertight by Open3D's strict test (the volume is taken as is, its sign
ignored); learned reconstruction is not compared.</p>
<div class="tablewrap"><table>
<thead><tr><th style='text-align:left'>object</th><th>mesh faces</th><th>FCP from mesh, error % (128 layers, 2048 atoms)</th><th>FCP from points, error % (128, 2048)</th>
<th>from mesh (128, 6144)</th><th>from points (128, 6144)</th><th>from mesh (256, 6144)</th><th>from points (256, 6144)</th>
<th>Poisson, true normals, error %</th><th>Poisson, estimated normals, error %</th><th>exact volume, &micro;s</th>
<th>FCP volume, &micro;s (128, 2048)</th><th>FCP encode from mesh, ms (CPU; 128, 2048)</th><th>FCP encode from points, ms (CPU; 128, 2048)</th>
<th>FCP encode from points, ms (CUDA; 128, 2048)</th><th>Poisson, true normals, ms</th><th>Poisson, estimated normals, ms</th></tr></thead>
<tbody>{orows}</tbody></table></div>
<figure><img src="sections/hull/volume_fig_renders.png" alt="Six objects: the original mesh, FCP decoded from the mesh and from 20k points, Poisson with true and with estimated normals, each with its volume error">
<figcaption>The surfaces behind the volumes: the original mesh, the FCP descriptor decoded (from the mesh and from 20k points;
the volume itself is computed from the atoms, not from this mesh), and Poisson's reconstruction with the true and with
estimated normals; under each, the volume error against the exact volume (FCP at 128 layers, 2048 atoms).</figcaption></figure>
"""
    (SECTION / "volume_section.html").write_text(html)
    note = ("\nvolume_*: ~/code/data/meshes/fcp_3010_hull/scripts/volume_study.py (steps in its docstring); "
                 "volume_data.csv: the first table; volume_objects.csv: per object (errors in %, times in ms / us as named).\n")
    readme = (SECTION / "README.txt").read_text()
    if note.strip() not in readme:
        (SECTION / "README.txt").write_text(readme.rstrip("\n") + "\n" + note)


if __name__ == "__main__":
    {"fcp": fcp_step, "poisson": poisson_step, "render": render_step, "section": section_step}.get(STEP, lambda: None)()