"""Pusher speed sweep with the CURRENT grip material + 500 mm blade.

Earlier speed tuning (pusher_diag*.py -> PUSH_SPEED=1.3) was measured while the blade was
still bound to the SLIPPERY DiverterMaterial (0.12/0.08); the grip fix (1.1/0.95) makes
those numbers stale. Also tests firing IMMEDIATELY at detection vs waiting to PUSH_X+0.08:
geometry says the item has only 0.5 m of blade (0.5 s at 1 m/s) and the wait burns 0.17 s
of it, so the wait may be why the stroke never completes on the item.

Outcome per trial: y_gain (needs > ~0.45 to reach the branch belt) and whether the item
survived at belt height or was ejected / fell through.
"""
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
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)

info = await plow_cell_9045.prepare(stage, belt_speed=1.0, script_control=True)
print("blade dims:", info["pusher_dims"], " seat:", info["pusher_seat"])

blade_prim = stage.GetPrimAtPath(_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_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])
r = bbc.ComputeWorldBound(stage.GetPrimAtPath("/World/Diverters/DiverterY_Split/Pusher/Geom")).ComputeAlignedRange()
BLADE_X0, BLADE_X1 = r.GetMin()[0], r.GetMax()[0]
SENSE_X = BLADE_X1          # leading (downstream-facing) edge
print(f"blade x[{BLADE_X0:+.3f}..{BLADE_X1:+.3f}]  sense at {SENSE_X:+.3f}  "
      f"contact window = {(BLADE_X1-BLADE_X0)/1.0:.3f} s at 1 m/s")

ITEM = "box_300x200x200"
ipath = "/World/Items/_pushprobe"

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" / f"{ITEM}.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.CreateSolverVelocityIterationCountAttr().Set(8)
    px.CreateMaxDepenetrationVelocityAttr().Set(C.MAX_DEPENETRATION)
    UsdGeom.Imageable(prim).MakeVisible()
    return RigidPrim(paths=[ipath])

print(f"\n{'speed':>6} {'wait':>6} {'y_gain':>8} {'final_x':>8} {'final_y':>8} {'final_z':>8}  verdict")
print("-" * 78)
results = []
for speed in (1.3, 1.7, 2.1, 2.6):
    for wait_to_center in (False, True):
        rp = spawn()
        blade_to(C.BLADE_HOME_Y)
        tl.play(); await app_utils.update_app_async(steps=8)
        # ride until the leading edge sees it
        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)
        if wait_to_center:
            for _ in range(60):
                if float(rp.get_world_poses()[0].numpy()[0][0]) <= C.PUSH_X + 0.08:
                    break
                await app_utils.update_app_async(steps=1)
        p0 = rp.get_world_poses()[0].numpy()[0].copy()
        a, b = C.BLADE_HOME_Y, 0.55
        dur = abs(b - a) / speed
        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
        for _ in range(60):          # let it settle / travel on
            await app_utils.update_app_async(steps=1)
        p = rp.get_world_poses()[0].numpy()[0]
        gain = float(p[1]) - float(p0[1])
        if float(p[2]) < 1.2:
            verdict = "FELL/LOST"
        elif abs(float(p[1])) > 3.0:
            verdict = "EJECTED"
        elif float(p[1]) > 0.45:
            verdict = "DELIVERED"
        else:
            verdict = "short - stayed on main belt"
        print(f"{speed:6.1f} {str(wait_to_center):>6} {gain:8.3f} {float(p[0]):8.3f} "
              f"{float(p[1]):8.3f} {float(p[2]):8.3f}  {verdict}")
        results.append((speed, wait_to_center, gain, verdict))
        tl.stop(); await app_utils.update_app_async(steps=6)
        blade_to(C.BLADE_HOME_Y)

good = [r for r in results if r[3] == "DELIVERED"]
print(f"\nDELIVERED configs: {[(s, w) for s, w, g, v in good]}")