"""fast_bounds2 against bounds.py (the method of the bounds page): agreement and time per object; out/bounds_fast.csv.

    env -u PYTHONPATH ~/code/fcp_devel/.venv/bin/python scripts/bounds_fast_check.py
"""
import sys
import time

import numpy as np

sys.argv = [sys.argv[0], "x"]
from bounds_fast import fast_bounds2, ring_coef_fast, ring_samples, candidates_par, _fourier, _obb6  # noqa: E402
from bounds import bounds  # noqa: E402
from hull import OBJECTS, OUT  # noqa: E402
import fcp  # noqa: E402
import pandas as pd  # noqa: E402
import trimesh  # noqa: E402


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


rows = []
for name, path in OBJECTS.items():
    m = trimesh.load(path, force="mesh")
    diag = float(np.linalg.norm(m.bounds[1] - m.bounds[0]))
    for layers, atoms in ((128, 2048), (128, 6144), (256, 6144)):
        d = fcp.encode(m, atoms=atoms, layers=layers, device="cpu")
        a0, o0, s0 = bounds(d)
        a1, o1, s1, npts = fast_bounds2(d)
        # the ring coefficients agree with the first fast version's (the Python loop)
        Ca, za = ring_coef_fast(d)
        Cb, zb = ring_samples(d, coef_only=True)
        assert np.allclose(Ca, Cb, atol=1e-9 * np.abs(Cb).max()) and np.array_equal(za, zb)
        C, z = Ca, za
        Fb = _fourier(d.K, 128)
        P = candidates_par(C, Fb, z, 64)
        W = d.pose.to_world(P)
        F = np.stack([np.eye(3), np.asarray(d.pose.rotation).T])
        r = dict(obj=name, layers=layers, atoms=atoms, points=npts,
                 now_ms=med(lambda: bounds(d), 31), fast_ms=med(lambda: fast_bounds2(d)),
                 fast_aabb_sphere_ms=med(lambda: fast_bounds2(d, obb=False)),
                 coef_ms=med(lambda: ring_coef_fast(d)), cand_ms=med(lambda: candidates_par(C, Fb, z, 64)),
                 obb_ms=med(lambda: _obb6(W, F)),
                 aabb_diff_pct=100 * float(np.abs(np.r_[a1[0] - a0[0], a1[1] - a0[1]]).max()) / diag,
                 obb_vol_ratio=float(np.prod(o1[2]) / np.prod(o0[2])), sphere_r_ratio=float(s1[1] / s0[1]))
        r["speedup"] = r["now_ms"] / r["fast_ms"]
        r["speedup_aabb_sphere"] = r["now_ms"] / r["fast_aabb_sphere_ms"]
        rows.append(r)
        print({k: (round(v, 4) if isinstance(v, float) else v) for k, v in r.items()}, flush=True)
d = pd.DataFrame(rows)
d.to_csv(OUT / "bounds_fast.csv", index=False)
print(d.groupby(["layers", "atoms"])[["now_ms", "fast_ms", "fast_aabb_sphere_ms", "speedup", "speedup_aabb_sphere", "coef_ms", "cand_ms", "obb_ms"]].median().round(2).to_string())
print(d.groupby(["layers", "atoms"]).agg(aabb=("aabb_diff_pct", "max"), obb_min=("obb_vol_ratio", "min"), obb_max=("obb_vol_ratio", "max"),
      s_min=("sphere_r_ratio", "min"), s_max=("sphere_r_ratio", "max"), fast_max=("fast_ms", "max"), sp_min=("speedup", "min")).round(4).to_string())