"""Mesh-mesh collision (do two placed solids overlap?) on random relative poses near contact: FCP straight from the two
descriptors (the ring vertices of each tested inside the other, CUDA and CPU numba, three sample strides) against
  standard: FCL (python-fcl, OBBRSS BVH, triangle-triangle), PyBullet on a V-HACD convex decomposition (GJK)
  SOTA:     Coal (the maintained successor of FCL/hpp-fcl), NVIDIA Warp (GPU: every triangle of A against the BVH of B,
            all poses in one launch, Moller triangle-triangle)
on the original meshes. Ground truth: surface intersection (Coal) or containment (a vertex of one inside the other) on
the original meshes; also on the decoded descriptors (the FCP geometry). Misses are measured: the penetration depth of
a missed collision and the separation of a false one, in % of the smaller object's box diagonal.
Writes ../report/collide.csv and ../report/collide_<pair>.npz.

    env -u PYTHONPATH .venv/bin/python scripts/collide_bench.py
"""
import csv
import os
import sys
import time
from pathlib import Path

import numpy as np
from scipy.spatial.transform import Rotation

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

import coal  # noqa: E402
import cupy as cp  # noqa: E402
import fcl  # noqa: E402
import open3d as o3d  # noqa: E402
import pybullet as pb  # noqa: E402
import warp as wp  # noqa: E402

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

wp.config.quiet = True
wp.init()
PAIRS = [("mug", "bowl"), ("dino", "horse"), ("shoe", "horse"), ("mug", "mug"), ("dress", "horse")]
if SET == "complex":
    PAIRS = [("spring", "spring"), ("knot", "knot"), ("sponge", "sponge"), ("gear", "gear"), ("pan_holes", "mug"), ("knot", "spring")]
if os.environ.get("FCP_PAIRS"):
    PAIRS = [tuple(p.split("-")) for p in os.environ["FCP_PAIRS"].split(",")]
K = int(os.environ.get("FCP_POSES", 2000))
_ls, _vs = max(1, LAYERS // 128), max(1, NT // 128)  # strides scaled so "1x1" is one sample per 128-layer / 128-vertex step
STRIDES = [(_ls * s, _vs * s) for s in (1, 2, 4)]  # (layer, vertex) strides of the FCP surface samples


@wp.kernel
def _wp_collide(meshB: wp.uint64, VA: wp.array(dtype=wp.vec3), FA: wp.array(dtype=wp.int32), VB: wp.array(dtype=wp.vec3),
                FB: wp.array(dtype=wp.int32), RT: wp.array(dtype=wp.mat33), t: wp.array(dtype=wp.vec3), hit: wp.array(dtype=wp.int32)):
    k, f = wp.tid()
    if hit[k] != 0:
        return
    a = RT[k] * (VA[FA[3 * f]] - t[k])
    b = RT[k] * (VA[FA[3 * f + 1]] - t[k])
    c = RT[k] * (VA[FA[3 * f + 2]] - t[k])
    lo = wp.min(wp.min(a, b), c)
    hi = wp.max(wp.max(a, b), c)
    q = wp.mesh_query_aabb(meshB, lo, hi)
    g = int(0)
    while wp.mesh_query_aabb_next(q, g):
        if wp.intersect_tri_tri(a, b, c, VB[FB[3 * g]], VB[FB[3 * g + 1]], VB[FB[3 * g + 2]]) != 0:
            hit[k] = 1
            return


def poses(mA, mB, rng):
    """B placed by x = R x_B + t: random rotation, its box centre at a random direction and distance u (rA + rB) from
    A's, u in [0.15, 1.0] (about half the poses collide)."""
    cA, cB = mA.bounds.mean(0), mB.bounds.mean(0)
    rA, rB = np.linalg.norm(mA.vertices - cA, axis=1).max(), np.linalg.norm(mB.vertices - cB, axis=1).max()
    R = Rotation.random(K, random_state=rng.integers(1 << 31)).as_matrix()
    v = rng.normal(size=(K, 3))
    v /= np.linalg.norm(v, axis=1, keepdims=True)
    u = rng.uniform(0.15, 1.0, K)
    t = cA + v * (u * (rA + rB))[:, None] - (R @ cB)
    return R, t


def coal_model(m):
    g = coal.BVHModelOBBRSS()
    g.beginModel(len(m.faces), len(m.vertices))
    g.addVertices(np.ascontiguousarray(m.vertices, float))
    g.addTriangles(np.ascontiguousarray(m.faces, np.int64))
    g.endModel()
    return g


def fcl_model(m):
    g = fcl.BVHModel()
    g.beginModel(len(m.vertices), len(m.faces))
    g.addSubModel(np.asarray(m.vertices, float), np.asarray(m.faces, np.int32))
    g.endModel()
    return g


def coal_hits(gA, gB, R, t, distance=False):
    req, out, dist = coal.CollisionRequest(), np.zeros(len(R), bool), np.zeros(len(R))
    I = coal.Transform3s()
    for k in range(len(R)):
        res = coal.CollisionResult()
        out[k] = coal.collide(gA, I, gB, coal.Transform3s(R[k], t[k]), req, res) > 0
        if distance and not out[k]:
            dr, dres = coal.DistanceRequest(), coal.DistanceResult()
            dist[k] = coal.distance(gA, I, gB, coal.Transform3s(R[k], t[k]), dr, dres)
    return (out, dist) if distance else out


def contains_any(mA, mB, R, t, sA, sB):
    """Containment: a vertex of A inside B or of B inside A (Open3D occupancy, 3-ray vote)."""
    xa = np.einsum("kji,kj->ki", R, mA.vertices[0] - t)  # A's vertex in B's frame: R^T (x - t)
    xb = (R @ mB.vertices[0]) + t
    ia = sB.compute_occupancy(o3d.core.Tensor(xa.astype(np.float32)), nsamples=3).numpy().astype(bool)
    ib = sA.compute_occupancy(o3d.core.Tensor(xb.astype(np.float32)), nsamples=3).numpy().astype(bool)
    return ia | ib


def scene(m):
    s = o3d.t.geometry.RaycastingScene()
    s.add_triangles(o3d.core.Tensor(np.asarray(m.vertices, np.float32)), o3d.core.Tensor(np.asarray(m.faces, np.uint32)))
    return s


def penetration(mA, mB, sA, sB, R, t, n=20000):
    """Deepest penetration (world units) of each pose: dense surface samples of each inside the other."""
    PA, PB = mA.sample(n), mB.sample(n)
    out = np.zeros(len(R))
    for k in range(len(R)):
        xa = (PA - t[k]) @ R[k]  # R^T (x - t)
        xb = PB @ R[k].T + t[k]
        # depth = distance of the samples inside the other (sign by a 3-ray vote: compute_signed_distance's one ray
        # misfires on perforated meshes and reported phantom penetrations of half the diagonal)
        dep = 0.0
        for X, s in ((xa, sB), (xb, sA)):
            T = o3d.core.Tensor(X.astype(np.float32))
            ins = s.compute_occupancy(T, nsamples=3).numpy() > 0
            if ins.any():
                dep = max(dep, float(s.compute_distance(T).numpy()[ins].max()))
        out[k] = dep
    return out


def bullet_setup(name, m, S):
    """V-HACD decomposition (cached) loaded as a compound of convex hulls, scaled by S."""
    src, dst = CACHE / f"{name}.obj", CACHE / f"{name}_vhacd.obj"
    t0 = time.perf_counter()
    if not dst.exists():
        m.export(src)
        pb.vhacd(str(src), str(dst), str(CACHE / f"{name}_vhacd.log"))
    t_vhacd = time.perf_counter() - t0
    shape = pb.createCollisionShape(pb.GEOM_MESH, fileName=str(dst), meshScale=[S] * 3)
    body = pb.createMultiBody(baseMass=0, baseCollisionShapeIndex=shape)
    return body, t_vhacd


def bullet_hits(bA, bB, R, t, S):
    out = np.zeros(len(R), bool)
    pb.resetBasePositionAndOrientation(bA, [0, 0, 0], [0, 0, 0, 1])
    for k in range(len(R)):
        q = Rotation.from_matrix(R[k]).as_quat()  # x, y, z, w (Bullet's order)
        pb.resetBasePositionAndOrientation(bB, (t[k] * S).tolist(), q.tolist())
        out[k] = len(pb.getClosestPoints(bA, bB, 0.0)) > 0
    return out


def main():
    pb.connect(pb.DIRECT)
    rows = []
    for nA, nB in PAIRS:
        tag = f"{nA}-{nB}"
        mA, mB = load_mesh(nA), load_mesh(nB)
        dA, dB = descriptor(nA), descriptor(nB)
        dmA, dmB = (fcp.decode(x, kind="mesh", resolution=min(x.layers, 256)) for x in (dA, dB))
        diag = min(np.linalg.norm(mA.extents), np.linalg.norm(mB.extents))
        rng = np.random.default_rng(7)
        R, t = poses(mA, mB, rng)
        sA, sB = scene(mA), scene(mB)
        # ground truths
        gA, gB = coal_model(mA), coal_model(mB)
        gt_surf, sep = coal_hits(gA, gB, R, t, distance=True)
        gt = gt_surf | contains_any(mA, mB, R, t, sA, sB)
        gt_dec = coal_hits(coal_model(dmA), coal_model(dmB), R, t) | contains_any(dmA, dmB, R, t, scene(dmA), scene(dmB))
        print(tag, "poses", K, "colliding", gt.mean(), "contained only", (gt & ~gt_surf).sum(), flush=True)
        common = dict(pair=tag, faces_A=len(mA.faces), faces_B=len(mB.faces), poses=K, collide_frac=float(gt.mean()))
        res = {}

        def record(meth, dev, build_s, per_pose_s, hits, extra=None):
            hits = np.asarray(hits, bool)
            fn, fp = gt & ~hits, ~gt & hits
            r = dict(common, method=meth, device=dev, build_ms=build_s * 1e3, us_per_pose=per_pose_s * 1e6,
                     acc=float((hits == gt).mean()), fn=int(fn.sum()), fp=int(fp.sum()),
                     acc_decoded=float((hits == gt_dec).mean()))
            r.update(extra or {})
            rows.append(r)
            res[meth] = hits
            print(f"  {meth:34s} build {r['build_ms']:9.1f} ms  {r['us_per_pose']:10.2f} us/pose  acc {r['acc']:.4f} "
                  f"(fn {r['fn']}, fp {r['fp']})  acc_dec {r['acc_decoded']:.4f}", flush=True)

        # --- FCP
        FCPQuery(dA, nt=NT)  # JIT warm-up (numba compile) outside the timed build
        t0 = time.perf_counter(); qA, qB = FCPQuery(dA, nt=NT), FCPQuery(dB, nt=NT); qA.gpu(); qB.gpu(); cuda_sync(); b_fcp = time.perf_counter() - t0
        maps = relative_maps(qA, qB, R, t)
        maps_g = [cp.asarray(x) for x in maps]
        for ls, vs in STRIDES:
            PA, PB = qA.surface_points(ls, vs), qB.surface_points(ls, vs)
            PAg, PBg = cp.asarray(PA), cp.asarray(PB)
            sec = timeit(lambda: collide_gpu(qA, qB, PAg, PBg, *maps_g), runs=5, sync=cuda_sync)
            hits = collide_gpu(qA, qB, PAg, PBg, *maps_g).get() > 0
            lab = f"{ls // _ls}x{vs // _vs}"
            record(f"FCP stride {lab} (CUDA)", "gpu", b_fcp, sec / K, hits, dict(samples=len(PA) + len(PB)))
            out = np.zeros(K, np.bool_)
            argsA, argsB = table_args(qA), table_args(qB)
            sec = timeit(lambda: collide_cpu_batch(PA, PB, *maps, argsA, argsB, out), runs=3)
            record(f"FCP stride {lab} (CPU)", "cpu", b_fcp, sec / K, out.copy(), dict(samples=len(PA) + len(PB)))
        # --- FCL
        t0 = time.perf_counter(); fA, fB = fcl_model(mA), fcl_model(mB); b = time.perf_counter() - t0
        oA, oB = fcl.CollisionObject(fA, fcl.Transform()), fcl.CollisionObject(fB, fcl.Transform())
        req = fcl.CollisionRequest()

        def run_fcl():
            h = np.zeros(K, bool)
            for k in range(K):
                oB.setTransform(fcl.Transform(R[k], t[k]))
                h[k] = fcl.collide(oA, oB, req, fcl.CollisionResult()) > 0
            return h

        t0 = time.perf_counter(); h = run_fcl(); sec = time.perf_counter() - t0
        record("FCL (BVH, tri-tri)", "cpu", b, sec / K, h)
        # --- Coal
        t0 = time.perf_counter(); cA, cB = coal_model(mA), coal_model(mB); b = time.perf_counter() - t0
        t0 = time.perf_counter(); h = coal_hits(cA, cB, R, t); sec = time.perf_counter() - t0
        record("Coal (BVH, tri-tri)", "cpu", b, sec / K, h)
        # --- PyBullet + V-HACD (scaled so the smaller object's diagonal is 10 units: Bullet's margins are absolute)
        S = 10.0 / diag
        bA, tvA = bullet_setup(nA, mA, S)
        bB, tvB = bullet_setup(nB, mB, S) if nB != nA else bullet_setup(nB, mB, S)
        t0 = time.perf_counter(); h = bullet_hits(bA, bB, R, t, S); sec = time.perf_counter() - t0
        record("PyBullet (V-HACD + GJK)", "cpu", tvA + tvB, sec / K, h)
        pb.removeBody(bA), pb.removeBody(bB)
        # --- Warp
        t0 = time.perf_counter()
        VA = wp.array(np.asarray(mA.vertices, np.float32), dtype=wp.vec3)
        FA = wp.array(np.asarray(mA.faces, np.int32).ravel(), dtype=wp.int32)
        VB = wp.array(np.asarray(mB.vertices, np.float32), dtype=wp.vec3)
        FB = wp.array(np.asarray(mB.faces, np.int32).ravel(), dtype=wp.int32)
        wmB = wp.Mesh(points=VB, indices=FB)
        wp.synchronize(); b = time.perf_counter() - t0
        RT = wp.array(np.transpose(R, (0, 2, 1)).astype(np.float32), dtype=wp.mat33)
        tw = wp.array(t.astype(np.float32), dtype=wp.vec3)
        hit = wp.zeros(K, dtype=wp.int32)

        def run_wp():
            hit.zero_()
            wp.launch(_wp_collide, dim=(K, len(mA.faces)), inputs=[wmB.id, VA, FA, VB, FB, RT, tw, hit])

        sec = timeit(run_wp, runs=3, sync=wp.synchronize)
        run_wp(); wp.synchronize()
        record("Warp (GPU BVH, tri-tri)", "gpu", b, sec / K, hit.numpy() > 0)
        # --- the misses of FCP (finest stride, CUDA): how deep / how far
        f = res["FCP stride 1x1 (CUDA)"]
        fn_idx, fp_idx = np.flatnonzero(gt & ~f), np.flatnonzero(~gt & f)
        depth = penetration(mA, mB, sA, sB, R[fn_idx], t[fn_idx]) if len(fn_idx) else np.zeros(0)
        np.savez_compressed(REPORT / f"collide{TAG}_{tag}.npz", R=R, t=t, gt=gt, gt_dec=gt_dec, sep=sep, diag=diag,
                            fn_idx=fn_idx, fn_depth=depth, fp_idx=fp_idx, fp_sep=sep[fp_idx],
                            **{k.replace(" ", "_"): v for k, v in res.items()})
        print(f"  FCP 1x1 misses: fn {len(fn_idx)} depth max {100 * depth.max() / diag if len(depth) else 0:.2f} % diag; "
              f"fp {len(fp_idx)} sep max {100 * sep[fp_idx].max() / diag if len(fp_idx) else 0:.2f} % diag", flush=True)
        with (REPORT / f"collide{TAG}.csv").open("w", newline="") as fh:
            keys = list(dict.fromkeys(k for r in rows for k in r))
            w = csv.DictWriter(fh, keys)
            w.writeheader()
            w.writerows(rows)


if __name__ == "__main__":
    main()