"""Re-time a sample of the grid on a quieter machine: the four objects that ran at the highest load, encodes 64/512,
128/2048, 256/6144, query settings (M, ring samples) (16, 64) and (64, 128); every task. out/tr_recheck.csv.

    env -u PYTHONPATH NUMBA_NUM_THREADS=1 OMP_NUM_THREADS=1 OPENBLAS_NUM_THREADS=1 MKL_NUM_THREADS=1 \
        ~/code/fcp_devel/.venv/bin/python scripts/tradeoffs_recheck.py
"""
import os
import sys

import numpy as np

sys.argv = [sys.argv[0], "x"]
from tradeoffs import med, OUT  # noqa: E402
from hull import OBJECTS  # noqa: E402
import pandas as pd  # noqa: E402
import trimesh  # noqa: E402
import fcp  # noqa: E402
from bounds_fast import fast_bounds3, hull_fast  # noqa: E402
from volume_atoms import volume_atoms_fast  # noqa: E402

assert os.environ.get("NUMBA_NUM_THREADS") == "1"
g = pd.read_csv(OUT / "tr_grid.csv").fillna({"frames": ""})
worst = g.groupby("obj").load.max().sort_values().index[-4:]
rows = []
for name in worst:
    m = trimesh.load(OBJECTS[name], force="mesh")
    for L, A in ((64, 512), (128, 2048), (256, 6144)):
        d = fcp.encode(m, atoms=A, layers=L, device="cpu")
        jobs = [("volume", 0, 0, "", lambda: volume_atoms_fast(d))]
        for M, nt in ((16, 64), (64, 128)):
            jobs += [("hull", M, nt, "", lambda M=M, nt=nt: hull_fast(d, M, nt)),
                     ("aabb", M, nt, "", lambda M=M, nt=nt: fast_bounds3(d, nt, M, which=("aabb",))),
                     ("sphere", M, nt, "", lambda M=M, nt=nt: fast_bounds3(d, nt, M, which=("sphere",))),
                     ("obb", M, nt, "own", lambda M=M, nt=nt: fast_bounds3(d, nt, M, frames="own", which=("obb",))),
                     ("obb", M, nt, "both", lambda M=M, nt=nt: fast_bounds3(d, nt, M, frames="both", which=("obb",)))]
        for task, M, nt, fr, f in jobs:
            q = g[(g.obj == name) & (g.layers == L) & (g.atoms == A) & (g.task == task) & (g.M == M) & (g.n_t == nt) & (g.frames == fr)].iloc[0]
            ms = med(f, 201 if task == "volume" else 21)
            rows.append(dict(obj=name, layers=L, atoms=A, task=task, M=M, n_t=nt, frames=fr, grid_ms=q.ms, grid_load=q.load,
                             ms=ms, load=os.getloadavg()[0], ratio=ms / q.ms))
    print(name, f"load {os.getloadavg()[0]:.2f}", flush=True)
d = pd.DataFrame(rows)
d.to_csv(OUT / "tr_recheck.csv", index=False)
print(d.groupby("task").ratio.describe()[["count", "min", "50%", "max"]].round(3).to_string())
print("overall median ratio (recheck / grid):", round(d.ratio.median(), 3), " load during recheck:", round(d.load.min(), 2), "-", round(d.load.max(), 2))
