"""Blade drive METHOD comparison - the speed sweep proved speed is not the variable.

A: write xformOp:translate      (what the pipeline does now) = a TELEPORT. PhysX sees no
   velocity; the item gets only a depenetration shove, so a FASTER blade pushes LESS -
   exactly what the sweep measured (1.3->99mm, 2.6->17mm). This is probe_push_physics.py's
   documented teleport signature.
B: RigidPrim.set_world_poses()  -> sets the KINEMATIC TARGET on the physics backend, so
   PhysX derives velocity = delta/dt and transfers real momentum.
C: dynamic blade + set_velocities() -> a genuine moving mass carrying momentum.
"""
import sys
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

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)
await plow_cell_9045.prepare(stage, belt_speed=1.0, script_control=True)

blade_prim = stage.GetPrimAtPath(_scene.BLADE)
blade_rp = RigidPrim(paths=[_scene.BLADE])
def _bop():
    for op in UsdGeom.Xformable(blade_prim).GetOrderedXformOps():
        if op.GetOpType() == UsdGeom.XformOp.TypeTranslate:
            return op
bop = _bop(); bbase = bop.Get()
def blade_xform_to(y):
    bop.Set(Gf.Vec3d(bbase[0], y - _scene.BLADE_PARENT_Y, bbase[2]))

bbc = UsdGeom.BBoxCache(Usd.TimeCode.Default(), [UsdGeom.Tokens.default_, UsdGeom.Tokens.render])
r0 = bbc.ComputeWorldBound(blade_prim).ComputeAlignedRange()
BLADE_WORLD_X = (r0.GetMin()[0] + r0.GetMax()[0]) / 2.0
BLADE_WORLD_Z = (r0.GetMin()[2] + r0.GetMax()[2]) / 2.0
SENSE_X = r0.GetMax()[0]
print(f"blade centre x={BLADE_WORLD_X:+.3f} z={BLADE_WORLD_Z:+.3f}, sense {SENSE_X:+.3f}")

ipath = "/World/Items/_pushprobe2"
def spawn():
    if stage.GetPrimAtPath(ipath).IsValid():
        stage.RemovePrim(ipath)
    prim = UsdGeom.Xform.Define(stage, ipath).GetPrim()
    prim.GetReferences().AddReference(str(C.ROOT / "assets" / "items" / "box_300x200x200.usd"))
    xf = UsdGeom.Xformable(prim); xf.ClearXformOpOrder()
    xf.AddTranslateOp(precision=UsdGeom.XformOp.PrecisionDouble).Set(Gf.Vec3d(-3.05, 0.0, C.BELT_Z + 0.05))
    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)
    px.CreateMaxDepenetrationVelocityAttr().Set(C.MAX_DEPENETRATION)
    UsdGeom.Imageable(prim).MakeVisible()
    return RigidPrim(paths=[ipath])

SPEED = 1.3
A, B = C.BLADE_HOME_Y, 0.55

async def ride_to_blade(rp):
    for _ in range(400):
        if float(rp.get_world_poses()[0].numpy()[0][0]) <= SENSE_X:
            return True
        await app_utils.update_app_async(steps=1)
    return False

print(f"\n{'method':>34} {'y_gain':>8} {'final_y':>8} {'final_z':>8}  verdict")
print("-" * 76)

for method in ("A: usd xform write (current)", "B: RigidPrim.set_world_poses",
               "C: dynamic + set_velocities"):
    # reset blade to kinematic home
    UsdPhysics.RigidBodyAPI(blade_prim).CreateKinematicEnabledAttr().Set(True)
    blade_xform_to(A)
    rp = spawn()
    tl.play(); await app_utils.update_app_async(steps=8)
    await ride_to_blade(rp)
    p0 = rp.get_world_poses()[0].numpy()[0].copy()

    dur = abs(B - A) / SPEED
    t0 = float(tl.get_current_time())
    if method.startswith("C"):
        UsdPhysics.RigidBodyAPI(blade_prim).CreateKinematicEnabledAttr().Set(False)
        UsdPhysics.MassAPI.Apply(blade_prim).CreateMassAttr().Set(200.0)
        PhysxSchema.PhysxRigidBodyAPI.Apply(blade_prim).CreateDisableGravityAttr().Set(True)
        await app_utils.update_app_async(steps=2)
    while True:
        t = float(tl.get_current_time()) - t0
        u = min(1.0, t / dur)
        y = A + (B - A) * u
        try:
            if method.startswith("A"):
                blade_xform_to(y)
            elif method.startswith("B"):
                blade_rp.set_world_poses(positions=np.array([[BLADE_WORLD_X, y, BLADE_WORLD_Z]]))
            else:
                vy = 0.0 if u >= 1.0 else SPEED
                blade_rp.set_velocities(np.array([[0.0, vy, 0.0, 0.0, 0.0, 0.0]]))
        except BaseException as exc:
            print(f"   {method}: drive call failed: {type(exc).__name__}")
            break
        await app_utils.update_app_async(steps=1)
        if u >= 1.0:
            break
    for _ in range(60):
        if method.startswith("C"):
            try:
                blade_rp.set_velocities(np.array([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]]))
            except BaseException:
                pass
        await app_utils.update_app_async(steps=1)

    p = rp.get_world_poses()[0].numpy()[0]
    gain = float(p[1]) - float(p0[1])
    verdict = ("FELL" if float(p[2]) < 1.2 else
               "EJECTED" if abs(float(p[1])) > 3.0 else
               "DELIVERED" if float(p[1]) > 0.45 else "short")
    print(f"{method:>34} {gain:8.3f} {float(p[1]):8.3f} {float(p[2]):8.3f}  {verdict}")

    tl.stop(); await app_utils.update_app_async(steps=6)
    UsdPhysics.RigidBodyAPI(blade_prim).CreateKinematicEnabledAttr().Set(True)
    blade_xform_to(A)