"""The error bound of the descriptor queries as a function of the setting: how far from the surface can an answer still be
wrong? Per setting (layers, N; the v2 encoder options: 2048 samples per ring, every island, cut at 1e-8 / 5 MB) and query
ring count NT (rule "nearest"):
  points:  over 1M near-surface (±2 % diag) and 1M uniform points per object, the largest distance to the true surface
           of a wrongly answered point (% of the diagonal), and counts beyond 0.1 / 0.25 / 0.5 / 1 / 2 %;
  collide: over 2000 poses near contact per pair, per sample stride, the deepest missed collision and the largest
           separation of a false one (% of the smaller diagonal). The true depth of every colliding pose and the
           separation of every free one are computed once per pair (cached).
Appends to ../report/bound_points.csv / bound_collide.csv.

    env -u PYTHONPATH .venv/bin/python scripts/bound.py points <layers> <N> [objects...]
    env -u PYTHONPATH .venv/bin/python scripts/bound.py collide <layers> <N> [pairs a-b ...]
    env -u PYTHONPATH .venv/bin/python scripts/bound.py truth [pairs a-b ...]
"""
import csv
import sys
import time
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).parent))
from common import CACHE, REPORT, cuda_sync, descriptor, load_mesh, timeit  # noqa: E402

import cupy as cp  # noqa: E402

from fcpq import FCPQuery, collide_gpu, relative_maps  # noqa: E402

OBJS = ["mug", "bowl", "dino", "horse", "shoe", "dress", "pan_holes", "knot", "spring", "sponge", "gear"]
PAIRS = ["mug-bowl", "dino-horse", "shoe-horse", "mug-mug", "dress-horse", "spring-spring", "knot-knot", "sponge-sponge",
         "gear-gear", "pan_holes-mug", "knot-spring"]
NTS = (128, 256, 512)
BANDS = (0.1, 0.25, 0.5, 1.0, 2.0)
SAMPLING = (1, 2, 4)  # collision samples every s/128 of the height and every s·NT/128-th ring vertex (as collide_bench)


def append(name, rows):
    out = REPORT / name
    new = not out.exists()
    with out.open("a", newline="") as f:
        w = csv.DictWriter(f, list(rows[0]))
        if new:
            w.writeheader()
        w.writerows(rows)


def points(layers, N, objs):
    from sweep import points as point_sets

    rows = []
    for name in objs:
        t0 = time.perf_counter()
        d = descriptor(name, layers, N)
        enc = time.perf_counter() - t0
        sets = point_sets(load_mesh(name))
        for nt in NTS:
            q = FCPQuery(d, nt=nt)
            row = dict(object=name, layers=layers, N=N, nt=nt, bytes=d.nbytes, parts=q.P, encode_or_load_s=enc)
            dmax, wrong_d = 0.0, []
            for k, (P, gt, dist) in sets.items():
                Pg = cp.asarray(d.pose.to_canonical(P).astype(np.float32))
                ins = q.inside_gpu(Pg, "nearest").get()
                row[f"{k}_ns"] = timeit(lambda: q.inside_gpu(Pg, "nearest"), runs=3, sync=cuda_sync) / len(P) * 1e9
                w = ins != gt
                row[f"{k}_wrong_pct"] = 100 * w.mean()
                wrong_d.append(dist[w])
                del Pg
            wd = np.concatenate(wrong_d)
            row["max_err"] = float(wd.max()) if len(wd) else 0.0
            row["p9999_err"] = float(np.percentile(wd, 99.99)) if len(wd) else 0.0
            for b in BANDS:
                row[f"n_beyond_{b}"] = int((wd > b).sum())
            rows.append(row)
            print(f"{name:9s} L{layers:4d} N{N:3d} NT{nt:4d} {d.nbytes / 1e6:6.3f} MB | wrong beyond: "
                  + " ".join(f"{b}%:{row[f'n_beyond_{b}']}" for b in BANDS) + f" | max {row['max_err']:.3f} % diag | "
                  f"{row['near_ns']:.1f} / {row['uniform_ns']:.1f} ns", flush=True)
            del q
            cp.get_default_memory_pool().free_all_blocks()
    append("bound_points.csv", rows)


def truth(pair):
    """Poses (collide_bench's: K 2000, seed 7), ground truth, true depth of colliding poses, separation of free ones."""
    import collide_bench as cb

    p = CACHE / f"truth_{pair}.npz"
    if p.exists():
        return dict(np.load(p))
    nA, nB = pair.split("-")
    mA, mB = load_mesh(nA), load_mesh(nB)
    diag = min(np.linalg.norm(mA.extents), np.linalg.norm(mB.extents))
    R, t = cb.poses(mA, mB, np.random.default_rng(7))
    sA, sB = cb.scene(mA), cb.scene(mB)
    surf, sep = cb.coal_hits(cb.coal_model(mA), cb.coal_model(mB), R, t, distance=True)
    gt = surf | cb.contains_any(mA, mB, R, t, sA, sB)
    depth = np.zeros(len(R))
    idx = np.flatnonzero(gt)
    depth[idx] = cb.penetration(mA, mB, sA, sB, R[idx], t[idx])
    out = dict(R=R, t=t, gt=gt, sep=sep, depth=depth, diag=diag)
    np.savez_compressed(p, **out)
    print(pair, f"truth: {gt.mean():.3f} collide; depth median {100 * np.median(depth[idx]) / diag:.3f} % diag", flush=True)
    return out


def collide(layers, N, pairs):
    rows = []
    for pair in pairs:
        T = truth(pair)
        nA, nB = pair.split("-")
        dA, dB = descriptor(nA, layers, N), descriptor(nB, layers, N)
        gt, diag = T["gt"], float(T["diag"])
        for nt in NTS:
            qA, qB = FCPQuery(dA, nt=nt), FCPQuery(dB, nt=nt)
            maps = [cp.asarray(x) for x in relative_maps(qA, qB, T["R"], T["t"])]
            for s in SAMPLING:
                ls, vs = max(1, layers // 128) * s, max(1, nt // 128) * s
                PA, PB = cp.asarray(qA.surface_points(ls, vs)), cp.asarray(qB.surface_points(ls, vs))
                hits = collide_gpu(qA, qB, PA, PB, *maps).get() > 0
                sec = timeit(lambda: collide_gpu(qA, qB, PA, PB, *maps), runs=3, sync=cuda_sync)
                fn, fp = gt & ~hits, ~gt & hits
                row = dict(pair=pair, layers=layers, N=N, nt=nt, sampling=s, samples=len(PA) + len(PB),
                           us_per_pose=sec / len(gt) * 1e6, acc=float((hits == gt).mean()), fn=int(fn.sum()), fp=int(fp.sum()),
                           fn_depth_max=float(100 * T["depth"][fn].max() / diag) if fn.any() else 0.0,
                           fp_sep_max=float(100 * T["sep"][fp].max() / diag) if fp.any() else 0.0)
                rows.append(row)
                print(f"{pair:14s} L{layers:4d} N{N:3d} NT{nt:4d} s{s}: {row['us_per_pose']:7.2f} us/pose acc {row['acc']:.4f} "
                      f"fn {row['fn']:3d} (deepest {row['fn_depth_max']:.3f} %) fp {row['fp']:3d} (widest {row['fp_sep_max']:.3f} %)",
                      flush=True)
            del qA, qB
            cp.get_default_memory_pool().free_all_blocks()
    append("bound_collide.csv", rows)


if __name__ == "__main__":
    mode = sys.argv[1]
    if mode == "points":
        points(int(sys.argv[2]), int(sys.argv[3]), sys.argv[4:] or OBJS)
    elif mode == "collide":
        collide(int(sys.argv[2]), int(sys.argv[3]), sys.argv[4:] or PAIRS)
    else:
        for p in sys.argv[2:] or PAIRS:
            truth(p)