"""Отработка пушера и плуга на новой сцене: поворот по классу и скольжение товара по лезвию.

Опирается на выясненное прогонами лент:
* узлы ConveyorBeltGraph нужно УДАЛЯТЬ - деактивация не мешает уже собранному графу
  обнулять surfaceVelocity на каждом тике;
* DiverterAnimGraph тоже удаляется: он переписывает цели приводов каждый тик и затирает
  всё, что задаёт Python;
* время берётся из таймлайна - заданная частота физики не применяется, фактический шаг
  83.33 мс (60 Гц), и от этого напрямую зависит вычисленная скорость.

Плуг ставится в положение ДО подхода товара - это и есть предпозиционирование по классу.
Скольжение по лезвию меряется как путь товара ВДОЛЬ кромки за время контакта: если товар
только отбрасывается, поперечная составляющая есть, а продольной нет.
"""
import sys, math
REPO = "/home/dasha/robozon-sorter"
if REPO not in sys.path:
    sys.path.insert(0, REPO)

import omni.usd, omni.timeline
import isaacsim.core.experimental.utils.app as app_utils
from pxr import Gf, Usd, UsdGeom, UsdPhysics, PhysxSchema, UsdShade
from isaacsim.core.experimental.prims import RigidPrim

from robozon_sorter import config as C
from robozon_sorter.sim import plow_cell

SPEED = 1.0
SCENE = f"{REPO}/scene/plow_cell_90_45_test.usd"
INTENT = {
    "/World/ConveyorTrack_05/Belt": (-1, 0, 0), "/World/ConveyorTrack/Belt": (-1, 0, 0),
    "/World/ConveyorTrack_02/Belt": (-1, 0, 0), "/World/ConveyorTrack_03/Belt": (-1, 0, 0),
    "/World/ConveyorTrack_04/Belt": (-1, 0, 0), "/World/ConveyorTrack_01/Belt": (-1, 0, 0),
    "/World/ConveyorTrack_06/Belt": (0, 1, 0),  "/World/ConveyorTrack_03/Belt_01": (0, 1, 0),
}

tl = omni.timeline.get_timeline_interface()
if tl.is_playing():
    tl.stop(); await app_utils.update_app_async(steps=10)
omni.usd.get_context().open_stage(SCENE)
await app_utils.update_app_async(steps=60)
stage = omni.usd.get_context().get_stage()

for path in [p.GetPath() for p in stage.Traverse()
             if "ConveyorBeltGraph" in p.GetName() or "DiverterAnimGraph" in p.GetName()]:
    stage.RemovePrim(path)
await app_utils.update_app_async(steps=20)

ps = next((p for p in stage.Traverse() if p.IsA(UsdPhysics.Scene)), None)
if ps is None:
    ps = UsdPhysics.Scene.Define(stage, "/World/PhysicsScene").GetPrim()
PhysxSchema.PhysxSceneAPI.Apply(ps).CreateEnableCCDAttr().Set(True)

GRIP = "/World/_BeltGrip"
g = stage.GetPrimAtPath(GRIP)
if not g.IsValid():
    g = stage.DefinePrim(GRIP, "Material")
pm = UsdPhysics.MaterialAPI.Apply(g)
pm.CreateStaticFrictionAttr().Set(1.1); pm.CreateDynamicFrictionAttr().Set(0.95)
pm.CreateRestitutionAttr().Set(0.02)
bb = UsdGeom.BBoxCache(Usd.TimeCode.Default(), [UsdGeom.Tokens.default_, UsdGeom.Tokens.render])
for path, intent in INTENT.items():
    pr = stage.GetPrimAtPath(path)
    if pr.IsValid():
        plow_cell.drive_belt(stage, path, intent, SPEED, grip_path=GRIP)
        PhysxSchema.PhysxSurfaceVelocityAPI(pr).CreateSurfaceVelocityEnabledAttr().Set(True)
TOP = bb.ComputeWorldBound(stage.GetPrimAtPath("/World/ConveyorTrack_04/Belt")
                           ).ComputeAlignedRange().GetMax()[2]
print(f"сцена готова, ленты на {SPEED} м/с, верх ленты z={TOP:.3f}")

# --- механизмы ------------------------------------------------------------------------
hinge = stage.GetPrimAtPath(C.PLOW_HINGE)
slide = stage.GetPrimAtPath("/World/Diverters/DiverterY_Split/PusherSlide")
print(f"плуг {C.PLOW_HINGE}: {hinge.IsValid()}   пушер PusherSlide: {slide.IsValid()}")
hdrive = UsdPhysics.DriveAPI(hinge, "angular")
pdrive = UsdPhysics.DriveAPI(slide, "linear")
arm = stage.GetPrimAtPath(C.PLOW_ARM)

def yaw_of(prim):
    M = UsdGeom.XformCache().GetLocalToWorldTransform(prim)
    d = M.TransformDir(Gf.Vec3d(1, 0, 0))
    return math.degrees(math.atan2(d[1], d[0]))

def spawn(name, x, y, size=0.05, mass=0.5):
    path = f"/World/_Goods/{name}"
    c = UsdGeom.Cube.Define(stage, path)
    c.CreateSizeAttr().Set(2.0)
    xf = UsdGeom.Xformable(c.GetPrim())
    xf.AddTranslateOp().Set(Gf.Vec3d(x, y, TOP + size + 0.005))
    xf.AddScaleOp().Set(Gf.Vec3f(size, size, size))
    p = c.GetPrim()
    UsdPhysics.RigidBodyAPI.Apply(p); UsdPhysics.CollisionAPI.Apply(p)
    UsdPhysics.MassAPI.Apply(p).CreateMassAttr().Set(mass)
    rb = PhysxSchema.PhysxRigidBodyAPI.Apply(p)
    rb.CreateEnableCCDAttr().Set(True); rb.CreateSolverPositionIterationCountAttr().Set(32)
    UsdShade.MaterialBindingAPI.Apply(p).Bind(
        UsdShade.Material(g), bindingStrength=UsdShade.Tokens.strongerThanDescendants,
        materialPurpose="physics")
    return path

async def run_case(label, plow_deg, start_x=-6.30, start_y=0.0, seconds=4.0):
    """поставить плуг ЗАРАНЕЕ, пустить товар, проследить его через плуг"""
    if stage.GetPrimAtPath("/World/_Goods").IsValid():
        stage.RemovePrim("/World/_Goods")
    stage.DefinePrim("/World/_Goods", "Xform")
    path = spawn(label, start_x, start_y)
    if hdrive:
        hdrive.GetTargetPositionAttr().Set(float(plow_deg))
        hdrive.CreateStiffnessAttr().Set(C.PLOW_STIFFNESS)
        hdrive.CreateDampingAttr().Set(C.PLOW_DAMPING)
        hdrive.CreateMaxForceAttr().Set(C.PLOW_MAX_FORCE)
    tl.play()
    await app_utils.update_app_async(steps=45)          # дать лезвию встать ДО подхода
    reached = yaw_of(arm) if arm.IsValid() else None
    rp = RigidPrim(paths=[path])
    T, P = [], []
    t_end = float(tl.get_current_time()) + seconds
    while float(tl.get_current_time()) < t_end:
        P.append(rp.get_world_poses()[0].numpy()[0].copy())
        T.append(float(tl.get_current_time()))
        await app_utils.update_app_async(steps=3)
    tl.stop(); await app_utils.update_app_async(steps=5)
    return label, plow_deg, reached, T, P

CASES = [("D_прямо", C.PLOW_PRESET["D"]), ("B_влево", C.PLOW_PRESET["B"]),
         ("C_вправо", C.PLOW_PRESET["C"]), ("B_широкий", C.PLOW_B_ANGLE)]
print(f"\nуглы из конфига: preset={C.PLOW_PRESET}  PLOW_B_ANGLE={C.PLOW_B_ANGLE}  "
      f"PLOW_X={C.PLOW_X}\n")
print(f"  {'случай':11s} {'цель°':>6s} {'факт°':>6s} {'смещ.Y,мм':>10s} {'скольж.,мм':>11s} "
      f"{'конец X,Y':>16s}  сторона")
print("  " + "-" * 84)
out = {}
for label, deg in CASES:
    label, deg, reached, T, P = await run_case(label, deg)
    if len(P) < 5:
        print(f"  {label:11s} нет данных"); continue
    y0 = float(P[0][1]); dy = float(P[-1][1]) - y0
    # скольжение по лезвию: путь ВДОЛЬ кромки за время контакта с зоной плуга
    slide_mm, prev = 0.0, None
    for p in P:
        if C.PLOW_SWEEP_X1 >= float(p[0]) >= C.PLOW_SWEEP_X0:
            if prev is not None:
                a = math.radians(reached if reached is not None else deg)
                ex, ey = math.cos(a), math.sin(a)       # направление кромки
                slide_mm += abs((float(p[0])-prev[0])*ex + (float(p[1])-prev[1])*ey) * 1000
            prev = (float(p[0]), float(p[1]))
    side = "+Y (B)" if dy > 0.05 else ("-Y (C)" if dy < -0.05 else "прямо")
    print(f"  {label:11s} {deg:6.1f} {(reached if reached is not None else float('nan')):6.1f} "
          f"{dy*1000:10.0f} {slide_mm:11.0f} ({P[-1][0]:+6.2f},{P[-1][1]:+6.2f})  {side}")
    out[label] = dict(target=deg, reached=reached, dy_mm=round(dy*1000),
                      slide_mm=round(slide_mm), end=[round(float(P[-1][0]), 2),
                                                     round(float(P[-1][1]), 2)], side=side)

# --- пушер ------------------------------------------------------------------------------
print("\nПУШЕР:")
if stage.GetPrimAtPath("/World/_Goods").IsValid():
    stage.RemovePrim("/World/_Goods")
stage.DefinePrim("/World/_Goods", "Xform")
gpath = spawn("push_D", -3.20, 0.0)
blade = stage.GetPrimAtPath("/World/Diverters/DiverterY_Split/Pusher")
tl.play(); await app_utils.update_app_async(steps=30)
rp = RigidPrim(paths=[gpath])
b0 = bb.ComputeWorldBound(blade).ComputeAlignedRange().GetMidpoint() if blade.IsValid() else None
fired = False
T, P, BY = [], [], []
t_end = float(tl.get_current_time()) + 5.0
while float(tl.get_current_time()) < t_end:
    p = rp.get_world_poses()[0].numpy()[0]
    P.append(p.copy()); T.append(float(tl.get_current_time()))
    if not fired and float(p[0]) <= C.PUSH_X + 0.10:
        if pdrive:
            pdrive.GetTargetPositionAttr().Set(0.45)     # выдвинуть
        fired = True
        print(f"   команда пушеру при x={float(p[0]):+.2f} (PUSH_X={C.PUSH_X})")
    if blade.IsValid():
        BY.append(float(bb.ComputeWorldBound(blade).ComputeAlignedRange().GetMidpoint()[1]))
    await app_utils.update_app_async(steps=3)
tl.stop(); await app_utils.update_app_async(steps=5)
if P:
    dy = float(P[-1][1]) - float(P[0][1])
    stroke = (max(BY) - min(BY)) * 1000 if BY else 0.0
    print(f"   ход лезвия пушера {stroke:.0f} мм")
    print(f"   товар: старт ({float(P[0][0]):+.2f},{float(P[0][1]):+.2f}) -> "
          f"конец ({float(P[-1][0]):+.2f},{float(P[-1][1]):+.2f}), смещение по Y {dy*1000:+.0f} мм")
    print(f"   {'ТОВАР УВЕДЁН НА ВЕТКУ' if dy > 0.15 else 'товар НЕ уведён'}")
    out["pusher"] = dict(stroke_mm=round(stroke), dy_mm=round(dy*1000))
globals()["PLOW_RESULT"] = out