#!/usr/bin/env python3
"""Close the two geometry faults that stop goods reaching a tray.

    /home/whatevenif/isaacsim/python.sh scripts/fix_plow_reach_and_trays.py

**1. Tray B is walled shut on the side goods arrive from.** Measured:

        B_W0 (near, y -2.55)  z 1.16 .. 2.00     <- 220 mm ABOVE the lane belt (1.7805)
        B_W1 (far,  y -3.35)  z 1.16 .. 1.70
        tray C, correctly built:
        C_W1 (near, y +1.88)  z 1.16 .. 1.70     <- 80 mm BELOW the belt, goods slide over
        C_W0 (far,  y +2.68)  z 1.16 .. 2.00

    B has its tall backboard on the near face instead of the far one, so the lane runs
    goods straight into a wall. The two heights are swapped to match C.

**2. The plow cannot reach the lane.** The arm is 600 mm on a hinge at the belt centre, so
    its tip reaches ``0.6 * sin(limit)``. At the authored +-35 deg that is 344 mm, while
    the lanes start at |y| = 450 mm: a 106 mm dead zone no command can cross. Goods are
    nudged to about y 0.15 and left on the line, which is exactly what every trace shows.

    Raising the joint limit to +-45 deg gives 424 mm of tip travel. That is still short of
    450 mm *at the centre of the item*, but an item is not a point: a 150 mm-wide box is
    carried by the lane once its near edge crosses, i.e. at a centre of about 375 mm, so
    42 deg (402 mm) delivers it with margin. Moving the lanes inboard instead was rejected -
    they would overlap the main belt, and two coincident belt colliders at the same height
    is its own failure.

Both edits are written back into scene/plow_cell.usd.
"""
from __future__ import annotations

import sys
from pathlib import Path

from pxr import Usd, UsdGeom, UsdPhysics

SCENE = Path(__file__).resolve().parent.parent / "scene" / "plow_cell.usd"
HINGE = "/World/Diverters/DiverterEnd/ArmHinge"
TRAYS = "/World/PlowContainers"

NEW_LIMIT = 45.0            # joint hard limit, degrees either side
NEAR_WALL_TOP = 1.70        # must sit below the lane belt at 1.7805
FAR_WALL_TOP = 2.00


def _range(stage, path):
    return UsdGeom.BBoxCache(0, ["default"]).ComputeWorldBound(
        stage.GetPrimAtPath(path)).ComputeAlignedRange()


def _set_top(stage, path, top):
    """scale a wall Cube about its base so its top lands at `top`"""
    prim = stage.GetPrimAtPath(path)
    r = _range(stage, path)
    z0, z1 = r.GetMin()[2], r.GetMax()[2]
    if abs(z1 - top) < 1e-4:
        return None
    want = max(top - z0, 0.02)
    xf = UsdGeom.Xformable(prim)
    for op in xf.GetOrderedXformOps():
        if op.GetOpType() == UsdGeom.XformOp.TypeScale:
            s = op.Get()
            op.Set(type(s)(s[0], s[1], s[2] * (want / max(z1 - z0, 1e-6))))
            break
    else:
        return None
    # keep the base where it was: scaling a centred cube moves it by half the delta
    for op in xf.GetOrderedXformOps():
        if op.GetOpType() == UsdGeom.XformOp.TypeTranslate:
            t = op.Get()
            op.Set(type(t)(t[0], t[1], t[2] + (want - (z1 - z0)) / 2.0))
            break
    return (round(z1, 3), round(_range(stage, path).GetMax()[2], 3))


def main():
    if not SCENE.exists():
        sys.exit(f"{SCENE} not found")
    stage = Usd.Stage.Open(str(SCENE))

    print("tray B - swap the tall backboard to the far face")
    for wall, top, tag in ((f"{TRAYS}/B_W0", NEAR_WALL_TOP, "near (goods arrive)"),
                           (f"{TRAYS}/B_W1", FAR_WALL_TOP, "far (backboard)")):
        if not stage.GetPrimAtPath(wall).IsValid():
            print(f"  {wall} missing"); continue
        changed = _set_top(stage, wall, top)
        print(f"  {wall.rsplit('/', 1)[1]:6s} {tag:22s} "
              + (f"top {changed[0]} -> {changed[1]}" if changed else "already correct"))

    print(f"plow hinge - raise the limit so the arm can reach the lane")
    hinge = stage.GetPrimAtPath(HINGE)
    if hinge.IsValid():
        j = UsdPhysics.RevoluteJoint(hinge)
        lo, hi = j.GetLowerLimitAttr().Get(), j.GetUpperLimitAttr().Get()
        j.GetLowerLimitAttr().Set(-NEW_LIMIT)
        j.GetUpperLimitAttr().Set(NEW_LIMIT)
        import math
        print(f"  limits {lo:+.0f}/{hi:+.0f} -> {-NEW_LIMIT:+.0f}/{NEW_LIMIT:+.0f}  "
              f"(tip reach {0.6 * math.sin(math.radians(abs(hi))):.3f} -> "
              f"{0.6 * math.sin(math.radians(NEW_LIMIT)):.3f} m, lane edge at 0.450)")
    else:
        print(f"  {HINGE} missing")

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


if __name__ == "__main__":
    main()