"""1) find the working RigidPrim velocity-set signature
   2) does Belt_01 (now aimed diagonally) actually carry an item into BinD?
   3) does a per-step 'carry assist' (matching the item's +Y velocity to the blade's)
      get a pushed item across, where the bare blade tops out at ~0.21 m?"""
import sys, inspect
REPO = "/home/dasha/robozon-sorter"
if REPO not in sys.path:
    sys.path.insert(0, REPO)
for _m in [k for k in list(sys.modules) if k.startswith("robozon_sorter")]:
    del sys.modules[_m]
import importlib; importlib.invalidate_caches()

import numpy as np
import omni.usd, omni.timeline
from pxr import Gf, Usd, UsdGeom, UsdPhysics, PhysxSchema
import isaacsim.core.experimental.utils.app as app_utils
from isaacsim.core.experimental.prims import RigidPrim
from robozon_sorter import config as C
from robozon_sorter.sim import scene as _scene, plow_cell_9045

print("set_velocities signature:", inspect.signature(RigidPrim.set_velocities))
print("get_velocities signature:", inspect.signature(RigidPrim.get_velocities))

stage = omni.usd.get_context().get_stage()
tl = omni.timeline.get_timeline_interface()
if tl.is_playing():
    tl.stop(); await app_utils.update_app_async(steps=10)
info = await plow_cell_9045.prepare(stage, belt_speed=1.0, script_control=True)
print("blade width:", info["pusher_dims"][0], "m")

ipath = "/World/Items/_imp"
def spawn(x, y, mesh="bag"):
    if stage.GetPrimAtPath(ipath).IsValid():
        stage.RemovePrim(ipath)
    prim = UsdGeom.Xform.Define(stage, ipath).GetPrim()
    prim.GetReferences().AddReference(str(C.ROOT / "assets" / "items" / f"{mesh}.usd"))
    xf = UsdGeom.Xformable(prim); xf.ClearXformOpOrder()
    xf.AddTranslateOp(precision=UsdGeom.XformOp.PrecisionDouble).Set(Gf.Vec3d(x, y, C.BELT_Z + 0.06))
    UsdPhysics.RigidBodyAPI.Apply(prim)
    UsdPhysics.RigidBodyAPI(prim).CreateKinematicEnabledAttr().Set(False)
    UsdPhysics.MassAPI.Apply(prim).CreateMassAttr().Set(0.6)
    px = PhysxSchema.PhysxRigidBodyAPI.Apply(prim)
    px.CreateEnableCCDAttr().Set(True)
    px.CreateSolverPositionIterationCountAttr().Set(24)
    UsdGeom.Imageable(prim).MakeVisible()
    return RigidPrim(paths=[ipath])

# ---------- 2) Belt_01 -> BinD ----------
print("\n--- item placed on Belt_01 at (-4.10,+0.70), does it reach BinD? ---")
rp = spawn(-4.10, 0.70)
tl.play(); await app_utils.update_app_async(steps=10)
for i in range(8):
    await app_utils.update_app_async(steps=45)
    p = rp.get_world_poses()[0].numpy()[0]
    print(f"   t~{(i+1)*45/60:4.1f}s x={float(p[0]):+.2f} y={float(p[1]):+.2f} z={float(p[2]):+.2f}")
p = rp.get_world_poses()[0].numpy()[0]
print("   IN BIN_D:", (-6.21 < float(p[0]) < -4.95) and (1.57 < float(p[1]) < 2.86))
tl.stop(); await app_utils.update_app_async(steps=6)

# ---------- 1)+3) velocity API and carry assist ----------
blade_prim = stage.GetPrimAtPath(_scene.BLADE)
for op in UsdGeom.Xformable(blade_prim).GetOrderedXformOps():
    if op.GetOpType() == UsdGeom.XformOp.TypeTranslate:
        bop = op; break
bbase = bop.Get()
def blade_to(y):
    bop.Set(Gf.Vec3d(bbase[0], y - _scene.BLADE_PARENT_Y, bbase[2]))

SENSE_X = C.PUSH_X + plow_cell_9045.PUSHER_X_MM / 2000.0
for mode in ("bare blade", "blade + carry assist"):
    blade_to(C.BLADE_HOME_Y)
    rp = spawn(-3.05, 0.0)
    tl.play(); await app_utils.update_app_async(steps=8)
    for _ in range(400):
        if float(rp.get_world_poses()[0].numpy()[0][0]) <= SENSE_X:
            break
        await app_utils.update_app_async(steps=1)
    p0 = rp.get_world_poses()[0].numpy()[0].copy()
    a, b, spd = C.BLADE_HOME_Y, 0.52, 1.3
    dur = abs(b - a) / spd
    t0 = float(tl.get_current_time()); err = None
    while True:
        u = min(1.0, (float(tl.get_current_time()) - t0) / dur)
        blade_to(a + (b - a) * u)
        if mode.endswith("assist") and u < 1.0:
            try:
                lin = rp.get_velocities()[0].numpy()[0]
                rp.set_velocities(np.array([[float(lin[0]), spd, float(lin[2])]]),
                                  np.array([[0.0, 0.0, 0.0]]))
            except BaseException as exc:
                err = f"{type(exc).__name__}: {exc}"
                break
        await app_utils.update_app_async(steps=1)
        if u >= 1.0:
            break
    if err:
        print(f"\n   {mode}: set_velocities FAILED -> {err}")
    else:
        for _ in range(90):
            await app_utils.update_app_async(steps=1)
        p = rp.get_world_poses()[0].numpy()[0]
        onbranch = float(p[1]) > 0.45
        print(f"\n   {mode}: dy={float(p[1])-float(p0[1]):+.3f}  final=({float(p[0]):+.2f},"
              f"{float(p[1]):+.2f},{float(p[2]):+.2f})  reached_branch={onbranch}")
    tl.stop(); await app_utils.update_app_async(steps=6)