#!/usr/bin/env python3
"""Rebuild the discharge as a Y fork with the plow at its apex.

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

The layout so far was a T: lane B perpendicular, lane C at 45 deg, meeting the main run at
different x, with the blade out in the middle of the belt trying to shove goods 400 mm
sideways onto them. Every failure this session came from that - the dead zone, the wedges,
goods stalling on the lip - because a push was being asked to do the job of a route.

A fork does not need the push. Both branches leave one apex, the blade sits in it as a
railway point, and goods **drive** into their branch:

    rest  (class C)  blade closes the B mouth  -> everything runs straight on to C
    B     arrives    blade swings over         -> the B mouth opens and takes it

Geometry, all from the apex at the downstream end of the main run:

    apex        x -7.00, y 0
    branch C    15 deg up   from the run   (travel -0.966, +0.259)
    branch B    35 deg down from the run   (travel -0.819, -0.574)

Each track is placed by measurement, not by assumption: the script reads where the belt slab
currently sits relative to its own prim origin, then sets the transform so the slab's near
end lands on the apex pointing along its branch. The trays follow their branches.
"""
from __future__ import annotations

import math
import sys
from pathlib import Path

from pxr import Gf, Usd, UsdGeom

SCENE = Path(__file__).resolve().parent.parent / "scene" / "plow_cell.usd"
APEX = (-7.00, 0.0)
TRAYS = "/World/PlowContainers"

# track prim -> (branch angle from the -X run, tray prefix, tray distance along the branch)
BRANCHES = {
    "/World/ConveyorTrack_01": dict(deg=+15.0, tray="C_", tray_at=2.30, name="C"),
    "/ConveyorTrack_01":       dict(deg=-35.0, tray="B_", tray_at=2.30, name="B"),
}


def _dir(deg):
    """world travel direction of a branch `deg` off the -X run"""
    a = math.radians(180.0 - deg)          # -X is 180 deg; +deg swings toward +Y
    return math.cos(a), math.sin(a)


def _belt_of(stage, track):
    for p in (f"{track}/Belt", f"{track}/Belt_01"):
        if stage.GetPrimAtPath(p).IsValid():
            return p
    return None


def _set_xform(stage, path, tx, ty, tz, rot_deg):
    xf = UsdGeom.Xformable(stage.GetPrimAtPath(path))
    xf.ClearXformOpOrder()
    xf.AddTranslateOp().Set(Gf.Vec3d(tx, ty, tz))
    xf.AddRotateZOp().Set(float(rot_deg))


def main():
    if not SCENE.exists():
        sys.exit(f"{SCENE} not found")
    stage = Usd.Stage.Open(str(SCENE))
    bb = UsdGeom.BBoxCache(0, ["default"])
    xc = UsdGeom.XformCache()

    for track, b in BRANCHES.items():
        prim = stage.GetPrimAtPath(track)
        if not prim.IsValid():
            print(f"  {track} missing"); continue
        belt = _belt_of(stage, track)
        if belt is None:
            print(f"  {track} has no Belt"); continue

        # where the slab sits now, relative to this track's own origin
        org = xc.GetLocalToWorldTransform(prim).ExtractTranslation()
        r = bb.ComputeWorldBound(stage.GetPrimAtPath(belt)).ComputeAlignedRange()
        span_x, span_y = r.GetMax()[0] - r.GetMin()[0], r.GetMax()[1] - r.GetMin()[1]
        length = max(span_x, span_y)
        top_z = r.GetMax()[2]
        # the slab's centre offset from the origin, in the track's own frame
        cx = (r.GetMin()[0] + r.GetMax()[0]) / 2.0 - org[0]
        cy = (r.GetMin()[1] + r.GetMax()[1]) / 2.0 - org[1]
        off = math.hypot(cx, cy)

        dx, dy = _dir(b["deg"])
        # put the slab centre half a length down the branch from the apex
        tx = APEX[0] + dx * (length / 2.0) - (dx * off - dx * off)
        ty = APEX[1] + dy * (length / 2.0)
        _set_xform(stage, track, tx - cx, ty - cy, org[2], 180.0 - b["deg"])

        r2 = bb.ComputeWorldBound(stage.GetPrimAtPath(belt)).ComputeAlignedRange()
        print(f"  branch {b['name']}  {b['deg']:+.0f} deg  dir ({dx:+.3f},{dy:+.3f})  "
              f"length {length:.2f} m")
        print(f"     slab now x[{r2.GetMin()[0]:+.2f},{r2.GetMax()[0]:+.2f}] "
              f"y[{r2.GetMin()[1]:+.2f},{r2.GetMax()[1]:+.2f}] top z={r2.GetMax()[2]:.3f}")

        # the tray rides to the end of its branch
        tex, tey = APEX[0] + dx * b["tray_at"], APEX[1] + dy * b["tray_at"]
        moved = 0
        for c in stage.GetPrimAtPath(TRAYS).GetChildren():
            if not c.GetName().startswith(b["tray"]):
                continue
            cr = bb.ComputeWorldBound(c).ComputeAlignedRange()
            ccx = (cr.GetMin()[0] + cr.GetMax()[0]) / 2.0
            ccy = (cr.GetMin()[1] + cr.GetMax()[1]) / 2.0
            cxf = UsdGeom.Xformable(c)
            for op in cxf.GetOrderedXformOps():
                if op.GetOpType() == UsdGeom.XformOp.TypeTranslate:
                    t = op.Get()
                    op.Set(type(t)(t[0] + (tex - ccx), t[1] + (tey - ccy), t[2]))
                    moved += 1
                    break
        print(f"     tray {b['name']} -> ({tex:+.2f},{tey:+.2f})  {moved} parts moved")

    stage.GetRootLayer().Save()
    print(f"saved {SCENE}")
    print("NOTE: the blade's rest position is now 'B closed', not 0 - re-measure which sign")
    print("      closes B before running a sort.")


if __name__ == "__main__":
    main()