#!/usr/bin/env python3
"""Lay the discharge out as a fork: C carries straight on, B branches, plow in the corner.

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

Design, from the sketch:

    main run ──────┬─────────────▶  lane C   (straight on, same line as the run)
                   │
                   └───▶ lane B   (branches away at 45 deg)
            plow sits in this corner

Class C needs no action - it runs straight through, and the blade at rest closes the B
mouth, so C is the default route and the blade only leans on it if it wanders. Class B is
the only case that actuates: the blade swings over, the B mouth opens, and the item drives
into its branch instead of being shoved across a belt.

**How the tracks are actually built** - this is what the first attempt got wrong. Each
ConveyorTrack carries `translate + orient(quaternion) + scale`; there is no rotateZ to
write, so clearing the op order and adding one silently produced a different transform. The
`Belt` child then sits at a fixed local offset of +1.0 along the track's local X, scaled by
the track's own X scale. So:

    belt centre = track origin + (local +X in world) * 1.0 * scale_x

Placing a branch therefore means: point the track's local X down the branch, and put the
track origin at the fork apex, which lands the belt centre one length-half down the branch.

Verification uses a **fresh** BBoxCache after every write. Reusing one is what made the
first attempt report "nothing moved" while the geometry underneath had in fact been
scattered.
"""
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 = Gf.Vec3d(-7.00, 0.0, 0.0)        # downstream end of the main run
TRAYS = "/World/PlowContainers"
PLOW = "/World/Diverters/DiverterEnd"

LANE_C = "/World/ConveyorTrack_01"       # straight on, 0 deg off the run
LANE_B = "/ConveyorTrack_01"             # branches 45 deg toward -Y
C_DEG, B_DEG = 0.0, -45.0
TRAY_AT = 2.60                           # how far down each branch its tray sits


def quat_z(deg):
    h = math.radians(deg) / 2.0
    return Gf.Quatd(math.cos(h), Gf.Vec3d(0, 0, math.sin(h)))


def world_dir(deg):
    """travel direction of a branch `deg` off the -X run"""
    a = math.radians(180.0 + deg)
    return Gf.Vec3d(math.cos(a), math.sin(a), 0.0)


def place_track(stage, path, deg):
    """point the track down its branch and hang its origin on the apex"""
    prim = stage.GetPrimAtPath(path)
    if not prim.IsValid():
        return False
    xf = UsdGeom.Xformable(prim)
    for op in xf.GetOrderedXformOps():
        n = op.GetOpName()
        if n.endswith("translate"):
            op.Set(APEX)
        elif n.endswith("orient"):
            op.Set(quat_z(180.0 + deg))      # local +X onto the branch direction
    return True


def move_group(stage, prefix, to_xy):
    """shift a tray so its centre lands on `to_xy`, keeping its parts together"""
    cache = UsdGeom.BBoxCache(0, ["default"])
    parts = [c for c in stage.GetPrimAtPath(TRAYS).GetChildren()
             if c.GetName().startswith(prefix)]
    if not parts:
        return 0
    xs, ys = [], []
    for c in parts:
        r = cache.ComputeWorldBound(c).ComputeAlignedRange()
        xs += [r.GetMin()[0], r.GetMax()[0]]
        ys += [r.GetMin()[1], r.GetMax()[1]]
    dx = to_xy[0] - (min(xs) + max(xs)) / 2.0
    dy = to_xy[1] - (min(ys) + max(ys)) / 2.0
    n = 0
    for c in parts:
        for op in UsdGeom.Xformable(c).GetOrderedXformOps():
            if op.GetOpName().endswith("translate"):
                t = op.Get()
                op.Set(type(t)(t[0] + dx, t[1] + dy, t[2]))
                n += 1
                break
    return n


def report(stage, path, label):
    """measure with a FRESH cache - a reused one reports the state before the write"""
    r = UsdGeom.BBoxCache(0, ["default"]).ComputeWorldBound(
        stage.GetPrimAtPath(path)).ComputeAlignedRange()
    if r.IsEmpty():
        print(f"  {label:22s} (empty)")
        return
    print(f"  {label:22s} x[{r.GetMin()[0]:+.2f},{r.GetMax()[0]:+.2f}] "
          f"y[{r.GetMin()[1]:+.2f},{r.GetMax()[1]:+.2f}] top z={r.GetMax()[2]:.3f}")


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

    for path, deg, tag in ((LANE_C, C_DEG, "C"), (LANE_B, B_DEG, "B")):
        if not place_track(stage, path, deg):
            print(f"  {path} missing"); continue
        d = world_dir(deg)
        tray = (APEX[0] + d[0] * TRAY_AT, APEX[1] + d[1] * TRAY_AT)
        moved = move_group(stage, f"{tag}_", tray)
        print(f"branch {tag}: {deg:+.0f} deg, dir ({d[0]:+.3f},{d[1]:+.3f}), "
              f"tray -> ({tray[0]:+.2f},{tray[1]:+.2f}) [{moved} parts]")

    # the plow sits in the corner between the two branches
    pxf = UsdGeom.Xformable(stage.GetPrimAtPath(PLOW))
    for op in pxf.GetOrderedXformOps():
        if op.GetOpName().endswith("translate"):
            t = op.Get()
            op.Set(Gf.Vec3d(APEX[0], APEX[1], t[2]))
            break

    stage.GetRootLayer().Save()
    print("\n--- measured after the write (fresh cache each time) ---")
    report(stage, "/World/ConveyorTrack_04/Belt", "main run")
    report(stage, f"{LANE_C}/Belt", "lane C (straight)")
    report(stage, f"{LANE_B}/Belt", "lane B (45 deg)")
    report(stage, f"{PLOW}/Arm", "plow arm")
    for tag in ("B", "C"):
        cache = UsdGeom.BBoxCache(0, ["default"])
        xs, ys = [], []
        for c in stage.GetPrimAtPath(TRAYS).GetChildren():
            if c.GetName().startswith(f"{tag}_"):
                r = cache.ComputeWorldBound(c).ComputeAlignedRange()
                xs += [r.GetMin()[0], r.GetMax()[0]]; ys += [r.GetMin()[1], r.GetMax()[1]]
        if xs:
            print(f"  tray {tag}                 centre "
                  f"({(min(xs)+max(xs))/2:+.2f},{(min(ys)+max(ys))/2:+.2f})")
    print(f"\nsaved {SCENE}")


if __name__ == "__main__":
    main()