"""1) Does Belt_01 actually CARRY an item, and does it carry it to BinD?
   2) How long does the pusher take to stroke out and return home?"""
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 omni.usd, omni.timeline
from pxr import Gf, Usd, UsdGeom, UsdPhysics, PhysxSchema, UsdShade
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)

bbc = UsdGeom.BBoxCache(Usd.TimeCode.Default(), [UsdGeom.Tokens.default_, UsdGeom.Tokens.render])
b01 = stage.GetPrimAtPath(_scene.BRANCH)
r = bbc.ComputeWorldBound(b01).ComputeAlignedRange()
print(f"Belt_01 bbox x[{r.GetMin()[0]:+.3f}..{r.GetMax()[0]:+.3f}] "
      f"y[{r.GetMin()[1]:+.3f}..{r.GetMax()[1]:+.3f}] top_z={r.GetMax()[2]:+.3f}")
sv = b01.GetAttribute("physxSurfaceVelocity:surfaceVelocity").Get()
en = b01.GetAttribute("physxSurfaceVelocity:surfaceVelocityEnabled")
xc = UsdGeom.XformCache()
wv = xc.GetLocalToWorldTransform(b01).TransformDir(Gf.Vec3d(*sv))
print(f"  local surfaceVelocity={sv}  enabled={en.Get() if en else None}")
print(f"  WORLD drive = ({wv[0]:+.3f},{wv[1]:+.3f},{wv[2]:+.3f})  |v|={wv.GetLength():.3f}")
api = UsdShade.MaterialBindingAPI(b01)
mat, _ = api.ComputeBoundMaterial(materialPurpose="physics")
if mat:
    m = UsdPhysics.MaterialAPI(mat.GetPrim())
    print(f"  friction static/dynamic = {m.GetStaticFrictionAttr().Get()}/{m.GetDynamicFrictionAttr().Get()}")
bd = bbc.ComputeWorldBound(stage.GetPrimAtPath("/World/SortingRig/BinD_Floor")).ComputeAlignedRange()
print(f"BinD floor x[{bd.GetMin()[0]:+.2f}..{bd.GetMax()[0]:+.2f}] y[{bd.GetMin()[1]:+.2f}..{bd.GetMax()[1]:+.2f}] "
      f"centre=({(bd.GetMin()[0]+bd.GetMax()[0])/2:+.2f},{(bd.GetMin()[1]+bd.GetMax()[1])/2:+.2f})")

ipath = "/World/Items/_branchprobe"
def spawn(x, y):
    if stage.GetPrimAtPath(ipath).IsValid():
        stage.RemovePrim(ipath)
    prim = UsdGeom.Xform.Define(stage, ipath).GetPrim()
    prim.GetReferences().AddReference(str(C.ROOT / "assets" / "items" / "bag.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])

print("\n--- placing an item DIRECTLY 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)
p0 = rp.get_world_poses()[0].numpy()[0].copy()
for i in range(9):
    await app_utils.update_app_async(steps=40)
    p = rp.get_world_poses()[0].numpy()[0]
    print(f"   t~{(i+1)*40/60:4.1f}s  x={float(p[0]):+.2f} y={float(p[1]):+.2f} z={float(p[2]):+.2f}")
pf = rp.get_world_poses()[0].numpy()[0]
in_bin = (-6.21 < float(pf[0]) < -4.95) and (1.57 < float(pf[1]) < 2.86)
print(f"   travelled ({float(pf[0])-float(p0[0]):+.2f},{float(pf[1])-float(p0[1]):+.2f})  IN BIN_D: {in_bin}")
tl.stop(); await app_utils.update_app_async(steps=6)

print("\n--- pusher stroke-out / return timing ---")
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]))
blade_to(C.BLADE_HOME_Y)
tl.play(); await app_utils.update_app_async(steps=5)
for label, a, b, spd in (("out ", C.BLADE_HOME_Y, 0.55, 1.3), ("back", 0.55, C.BLADE_HOME_Y, 1.3),
                          ("back", 0.55, C.BLADE_HOME_Y, 2.5)):
    blade_to(a); await app_utils.update_app_async(steps=3)
    dur = abs(b - a) / spd
    t0 = float(tl.get_current_time())
    while True:
        u = min(1.0, (float(tl.get_current_time()) - t0) / dur)
        blade_to(a + (b - a) * u)
        await app_utils.update_app_async(steps=1)
        if u >= 1.0:
            break
    print(f"   {label} @ {spd} m/s : {float(tl.get_current_time())-t0:.3f} s "
          f"(stroke {abs(b-a):.2f} m)")
tl.stop()