"""Door robot arm: where the tool passes over the wrist block on the doors' real path (stage 3 fix round, CRITIQUE_2 must 1).

Reads the rigged file: for both versions and both doors, at every frame 1-120, every vertex of the tool (version A: the
hub and the eight finger links as posed by the rig, grip closed; version B: the neck and plate) is taken into the
forearm's own frame (origin J4, the wrist centre at x 1.034). Per band of |y| (the central wedge, the cheeks, the J5
caps) and per 4 mm of x, the lowest point of the tool above the wrist head is kept. The wrist block's sloped top is cut
back only where this envelope (less a clearance) comes lower than the slope.

    blender -b --factory-startup --python build/dra_wrist_envelope.py -- <file.blend>   -> build/wrist_envelope.json
                                                                                          and wrist_envelope_grid.json
The grid file keeps the lowest tool point per 4 mm of x and 5 mm of |y| (so the cheek's top can be cut back only
where the tool really passes, and its outer edge kept).
"""
import os
import sys
import json
import numpy as np
import bpy

HERE = os.path.dirname(os.path.abspath(__file__))
ARGS = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
SUB = int(ARGS[ARGS.index("--sub") + 1]) if "--sub" in ARGS else 1      # steps per frame (4: every quarter frame)
ARGS = [a for i, a in enumerate(ARGS) if a != "--sub" and (i == 0 or ARGS[i - 1] != "--sub")]
SWEEP = "--sweep" in ARGS
ARGS = [a for a in ARGS if a != "--sweep"]
path = ARGS[0] if ARGS else os.path.join(os.path.dirname(HERE), "DoorRobotArm.blend")
bpy.ops.wm.open_mainfile(filepath=path)
if SWEEP:
    # the geo_sweep's own poses (dra_sweep._j5_coupled): J5 up to MARGIN past the path's ends with J6 as the path has
    # it there, J6 Q6_BAND either side of the path at every other path point; the showroom arms' rest geometry
    # (fingers open). -> build/wrist_envelope_sweep.json (same grid as below)
    sys.path.insert(0, HERE)
    import math
    import dra_sweep as sw
    import dra_ik
    D = sw.design()
    dps = sw.door_paths()
    sc = bpy.data.scenes["DoorArm | Showroom"]
    if bpy.context.window_manager.windows:
        bpy.context.window_manager.windows[0].scene = sc
    dg = sc.view_layers[0].depsgraph
    dg.update()
    X0, X1, DX = 0.50, 1.20, 0.004
    nb = int((X1 - X0) / DX) + 1
    Y0, DY, ny = 0.0, 0.005, 50
    res = {}
    for ver, tools in (("A", ["tool A hub"] + ["finger %s %s" % (w, q) for w in ("top", "left", "bottom", "right")
                                              for q in ("proximal", "distal")]), ("B", ["tool B"])):
        p = dra_ik.path(D, ver, dps)
        pts = sorted({(round(math.degrees(q["J5"]), 1), round(math.degrees(q["J6"]), 1)) for _, _, q in p})
        q5s = [a for a, _ in pts]
        deep = [pt for pt in pts if pt[0] <= min(q5s) + 0.5]
        shallow = [pt for pt in pts if pt[0] >= max(q5s) - 0.5]
        tests = set()
        for (a, b) in pts[::2] + deep + shallow:
            for db in (-sw.Q6_BAND, 0.0, sw.Q6_BAND):
                tests.add((a, round(b + db, 1)))
        for k in (1, 2, 3):
            e = sw.MARGIN * k / 3
            for (a, b) in deep:
                tests.add((round(a - e, 1), b))
            for (a, b) in shallow:
                tests.add((round(a + e, 1), b))
        prefix = "Showroom | arm %s | " % ver
        R0 = bpy.data.objects["Showroom | arm " + ver].matrix_world.copy()
        objs = {o["link"]: o for o in sc.objects if o.name.startswith(prefix) and o.type == 'MESH' and "link" in o}
        Finv = np.array(objs["forearm"].matrix_world.inverted())
        W = {}
        for t in tools:
            v, _, _ = sw.world_geo(objs[t], dg)
            W[t] = np.array([tuple(x) for x in v], dtype=np.float64)
        grid = np.full((nb, ny), 9.0)
        for (a, b) in sorted(tests):
            mats = sw._wrist_mats(R0, a, b, tools)
            for t in tools:
                M = Finv @ np.array(mats[t])
                q = W[t] @ M[:3, :3].T + M[:3, 3]
                ay = np.abs(q[:, 1])
                g = (q[:, 2] > 0.04) & (q[:, 0] >= X0) & (q[:, 0] < X1) & (ay < Y0 + DY * ny)
                if g.any():
                    np.minimum.at(grid, (((q[g, 0] - X0) / DX).astype(int), ((ay[g] - Y0) / DY).astype(int)), q[g, 2])
        res[ver] = {"x0": X0, "dx": DX, "y0": Y0, "dy": DY, "poses": len(tests),
                    "z": [[round(float(v), 4) if v < 9 else None for v in row] for row in grid]}
        print("[dra] sweep envelope", ver, len(tests), "poses", flush=True)
    with open(os.path.join(HERE, "wrist_envelope_sweep.json"), "w") as fh:
        json.dump(res, fh)
    print("[dra] wrist sweep envelope written", flush=True)
    raise SystemExit(0)
sf = bpy.data.scenes["DoorArm | Ship fit"]
wm = bpy.context.window_manager
if wm.windows:
    wm.windows[0].scene = sf
ctl = bpy.data.objects["DoorArm | CONTROLS"]
BANDS = {"centre": (0.0, 0.095), "cheek": (0.105, 0.180), "cap": (0.180, 0.230)}
X0, X1, DX = 0.50, 1.20, 0.004
nb = int((X1 - X0) / DX) + 1
Y0, DY, ny = 0.0, 0.005, 50                 # the fine grid: |y| in 5 mm steps up to 0.25
out = {}
grid_out = {}
for ver, tools in (("A", ["tool A hub"] + ["finger %s %s" % (w, p) for w in ("top", "left", "bottom", "right")
                                          for p in ("proximal", "distal")]), ("B", ["tool B"])):
    ctl["version"] = 0 if ver == "A" else 1
    ctl.update_tag()
    env = {b: np.full(nb, 9.0) for b in BANDS}
    grid = np.full((nb, ny), 9.0)
    who = {b: [None] * nb for b in BANDS}
    for key in ("AFT", "FORE"):
        pre = "DoorArm_%s_%s" % (key, ver)
        fa = bpy.data.objects["%s | forearm" % pre]
        obs = [bpy.data.objects["%s | %s" % (pre, t)] for t in tools]
        cos = {}
        for o in obs:
            co = np.empty(len(o.data.vertices) * 3, np.float64)
            o.data.vertices.foreach_get("co", co)
            cos[o.name] = co.reshape(-1, 3)
        for fs in range(SUB, 120 * SUB + 1):
            f = fs / SUB
            sf.frame_set(int(f), subframe=f - int(f))
            dg = bpy.context.evaluated_depsgraph_get()
            Finv = np.array(fa.evaluated_get(dg).matrix_world.inverted())
            for o in obs:
                M = Finv @ np.array(o.evaluated_get(dg).matrix_world)
                p = cos[o.name] @ M[:3, :3].T + M[:3, 3]
                ay = np.abs(p[:, 1])
                g = (p[:, 2] > 0.04) & (p[:, 0] >= X0) & (p[:, 0] < X1) & (ay < Y0 + DY * ny)
                if g.any():
                    gx = ((p[g, 0] - X0) / DX).astype(int)
                    gy = ((ay[g] - Y0) / DY).astype(int)
                    np.minimum.at(grid, (gx, gy), p[g, 2])
                for b, (y0, y1) in BANDS.items():
                    m = (ay >= y0) & (ay < y1) & (p[:, 2] > 0.04) & (p[:, 0] >= X0) & (p[:, 0] < X1)
                    if not m.any():
                        continue
                    q = p[m]
                    ix = ((q[:, 0] - X0) / DX).astype(int)
                    for i, z in zip(ix, q[:, 2]):
                        if z < env[b][i]:
                            env[b][i] = z
                            who[b][i] = [key, round(f, 2), o.name.split(" | ")[-1]]
    out[ver] = {b: [[round(X0 + i * DX, 3), round(float(env[b][i]), 4), who[b][i]] for i in range(nb) if env[b][i] < 9.0]
                for b in BANDS}
    grid_out[ver] = {"x0": X0, "dx": DX, "y0": Y0, "dy": DY,
                     "z": [[round(float(v), 4) if v < 9 else None for v in row] for row in grid]}
    print("[dra] envelope", ver, {b: (round(float(env[b].min()), 3) if (env[b] < 9).any() else None) for b in BANDS}, flush=True)
ctl["version"] = 0
ctl.update_tag()
with open(os.path.join(HERE, "wrist_envelope.json"), "w") as fh:
    json.dump(out, fh, indent=0)
with open(os.path.join(HERE, "wrist_envelope_grid.json"), "w") as fh:
    json.dump(grid_out, fh)
print("[dra] wrist envelope written", flush=True)