#!/usr/bin/env python3
"""Bring both plow lanes inboard so the blade can actually reach them.

    /home/whatevenif/isaacsim/python.sh scripts/move_lanes_inboard.py [--shift 0.07]

Measured over the 25-object run: **no item ever crossed y = 0.39**, while the lanes start
at |y| = 0.45. 16 of 25 came to rest against the blade still on the belt. 0.39 is not a
coincidence - it is where the blade tip is: a 600 mm arm at 42 deg reaches 0.6*sin42 =
0.402 m, and the item is pushed to the tip and no further, because that is where the blade
ends. Raising the angle cannot close it either: at the joint's 45 deg limit the tip reaches
0.424, still short.

So the lane comes to the blade. Each lane moves 70 mm toward the centreline, putting its
near edge at |y| = 0.38 - inside the tip's reach with ~20 mm to spare.

**Its tray moves with it.** Moving the lane alone would widen the gap between the lane end
and the tray wall from 100 mm to 170 mm, and goods would fall short onto the floor instead
of into the tray. Lane and tray are one assembly and are shifted by the same vector.

Known cost: the lane now overlaps the carrying belt by 70 mm (the belt is +-0.45 wide).
Two belt colliders share that strip at the same height, with different surface velocities.
That strip is exactly the hand-over region, so goods there being pulled by both is the
intended behaviour rather than a defect - but it is the thing to look at first if items
start jittering at the lane entry.
"""
from __future__ import annotations

import argparse
import sys
from pathlib import Path

from pxr import Gf, Usd, UsdGeom

SCENE = Path(__file__).resolve().parent.parent / "scene" / "plow_cell.usd"

# prim -> how far to move it in +y (toward the centreline from its own side)
GROUPS = {
    "B": dict(shift=+1.0, prims=["/ConveyorTrack_01", "/World/PlowCornerDeck_B"],
              tray_prefix="B_"),
    "C": dict(shift=-1.0, prims=["/World/ConveyorTrack_01", "/World/PlowCornerDeck_C",
                                 "/World/PlowTransition_C"],
              tray_prefix="C_"),
}
TRAYS = "/World/PlowContainers"


def _shift_y(stage, path, dy):
    prim = stage.GetPrimAtPath(path)
    if not prim.IsValid():
        return False
    xf = UsdGeom.Xformable(prim)
    for op in xf.GetOrderedXformOps():
        if op.GetOpType() == UsdGeom.XformOp.TypeTranslate:
            t = op.Get()
            op.Set(type(t)(t[0], t[1] + dy, t[2]))
            return True
        if op.GetOpType() == UsdGeom.XformOp.TypeTransform:
            M = Gf.Matrix4d(op.Get())
            tr = M.ExtractTranslation()
            M.SetTranslateOnly(Gf.Vec3d(tr[0], tr[1] + dy, tr[2]))
            op.Set(M)
            return True
    xf.AddTranslateOp().Set(Gf.Vec3d(0.0, dy, 0.0))
    return True


def _edge(stage, path):
    r = UsdGeom.BBoxCache(0, ["default"]).ComputeWorldBound(
        stage.GetPrimAtPath(path)).ComputeAlignedRange()
    return None if r.IsEmpty() else (r.GetMin()[1], r.GetMax()[1])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--shift", type=float, default=0.07, help="metres toward the centreline")
    args = ap.parse_args()
    if not SCENE.exists():
        sys.exit(f"{SCENE} not found")
    stage = Usd.Stage.Open(str(SCENE))

    for tag, g in GROUPS.items():
        dy = g["shift"] * args.shift
        before = _edge(stage, g["prims"][0])
        moved = [p for p in g["prims"] if _shift_y(stage, p, dy)]
        trays = [c.GetPath().pathString
                 for c in stage.GetPrimAtPath(TRAYS).GetChildren()
                 if c.GetName().startswith(g["tray_prefix"])]
        moved += [p for p in trays if _shift_y(stage, p, dy)]
        after = _edge(stage, g["prims"][0])
        near_before = min(abs(v) for v in before) if before else float("nan")
        near_after = min(abs(v) for v in after) if after else float("nan")
        print(f"lane {tag}: dy {dy:+.3f} m, {len(moved)} prims "
              f"({len(trays)} of them tray parts)")
        print(f"          near edge |y| {near_before:.3f} -> {near_after:.3f}  "
              f"(blade tip reaches 0.402)")

    stage.GetRootLayer().Save()
    print(f"saved {SCENE}")


if __name__ == "__main__":
    main()