#!/usr/bin/env python3
"""Carry the driven surface inboard, to where the plow can actually deliver.

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

Measured with one item and a full 42 deg sweep: the blade imparts a real push (item picks up
0.59 m/s and travels 139 mm sideways) but leaves it at **y = 0.261**, and every driven
surface past the plow starts at **|y| = 0.380**:

    belt  ConveyorTrack_04   x -7.00..-6.00   y -0.450..+0.450
    deck  PlowTransition_C   x -7.00..-6.00   y +0.380..+0.698
    lane  C                  x -8.12..-6.39   y +0.380..+2.112
    lane  B                  x -7.03..-6.58   y -2.379..-0.380   (no transition deck at all)

So there is a 119 mm band where a swept item sits on the very lip of the main belt with
nothing driving it toward its lane. That is the gap the goods die in - not a hole they fall
through, a strip with no traction, right where the blade lets go of them.

This closes it from the inside: each transition plate is brought in to |y| = 0.20, well
short of where the blade releases, and B gets the plate it never had. `plow_sort.DECK_DIR`
already drives both toward their lanes, so an item landing here is carried on instead of
stopping.

Plates are static Cubes with collision, coplanar with the belt at z 1.7805, and are made
kinematic + surface-driven at load time like every other deck.
"""
from __future__ import annotations

import sys
from pathlib import Path

from pxr import Gf, Sdf, Usd, UsdGeom, UsdPhysics

SCENE = Path(__file__).resolve().parent.parent / "scene" / "plow_cell.usd"
TOP_Z = 1.7805
THICK = 0.01
INBOARD = 0.20          # how far in the driven surface now reaches

# Lane B is perpendicular, so a straight strip meets it flush.
PLATES_STRAIGHT = {"PlowTransition_B": (-7.03, -6.45, -0.380)}

# Lane C is laid at 45 deg. Its bounding box, x -8.123..-6.391 by y 0.380..2.112, is the
# AABB of a rotated rectangle and describes a footprint the belt does not have: the near
# side is a single CORNER at (-7.257, 0.380), and the edge runs away from it at 45 deg,
#
#     y = x + 7.637
#
# so at x -7.03 the lane really starts at y 0.607, and at x -6.39 at y 1.247 - not at 0.380
# anywhere except that one corner. A straight plate ending at y 0.380 therefore leaves a
# widening wedge of open air, which is the dark triangle in the viewport and where
# `bolts_cluster` fell through after the blade had successfully pushed it to y +0.425.
#
# The C plate is a trapezoid instead: inboard edge at |y| = INBOARD, outer edge ON the
# lane's diagonal.
LANE_C_EDGE = lambda x: x + 7.637


def _plate(stage, name, x0, x1, y_in, y_out):
    path = f"/World/{name}"
    prim = stage.GetPrimAtPath(path)
    if prim.IsValid():
        stage.RemovePrim(path)
    cube = UsdGeom.Cube.Define(stage, path)
    cube.CreateSizeAttr().Set(2.0)          # size 2 so the scale op IS the half-extent
    cx, cy = (x0 + x1) / 2.0, (y_in + y_out) / 2.0
    hx, hy = abs(x1 - x0) / 2.0, abs(y_out - y_in) / 2.0
    xf = UsdGeom.Xformable(cube.GetPrim())
    xf.ClearXformOpOrder()
    xf.AddTranslateOp().Set(Gf.Vec3d(cx, cy, TOP_Z - THICK))
    xf.AddScaleOp().Set(Gf.Vec3f(hx, hy, THICK))
    cube.CreateDisplayColorAttr().Set([Gf.Vec3f(0.30, 0.31, 0.33)])
    UsdPhysics.CollisionAPI.Apply(cube.GetPrim())
    return path, (round(cx - hx, 3), round(cx + hx, 3), round(cy - hy, 3), round(cy + hy, 3))


def _trapezoid(stage, name, quad):
    """a thin prism whose top face is the given 4 corners, coplanar with the belt.

    A Cube cannot do this - the plate has to follow a 45 deg edge, so it is authored as an
    explicit mesh. Given thickness rather than left as a zero-height quad: a flat sheet is a
    poor collider and goods catch on its rim.
    """
    path = f"/World/{name}"
    if stage.GetPrimAtPath(path).IsValid():
        stage.RemovePrim(path)
    mesh = UsdGeom.Mesh.Define(stage, path)
    top = [Gf.Vec3f(x, y, TOP_Z) for x, y in quad]
    bot = [Gf.Vec3f(x, y, TOP_Z - THICK) for x, y in quad]
    pts = top + bot
    mesh.CreatePointsAttr().Set(pts)
    faces, counts = [], []
    faces += [0, 1, 2, 3]; counts.append(4)                  # top
    faces += [7, 6, 5, 4]; counts.append(4)                  # bottom
    for i in range(4):                                       # sides
        j = (i + 1) % 4
        faces += [i, 4 + i, 4 + j, j]; counts.append(4)
    mesh.CreateFaceVertexIndicesAttr().Set(faces)
    mesh.CreateFaceVertexCountsAttr().Set(counts)
    xs = [p[0] for p in pts]; ys = [p[1] for p in pts]; zs = [p[2] for p in pts]
    mesh.CreateExtentAttr().Set([Gf.Vec3f(min(xs), min(ys), min(zs)),
                                 Gf.Vec3f(max(xs), max(ys), max(zs))])
    mesh.CreateDisplayColorAttr().Set([Gf.Vec3f(0.30, 0.31, 0.33)])
    mesh.CreateSubdivisionSchemeAttr().Set("none")
    UsdPhysics.CollisionAPI.Apply(mesh.GetPrim())
    UsdPhysics.MeshCollisionAPI.Apply(mesh.GetPrim()).CreateApproximationAttr().Set("convexHull")
    return path


def main():
    if not SCENE.exists():
        sys.exit(f"{SCENE} not found")
    stage = Usd.Stage.Open(str(SCENE))
    for name, (x0, x1, y_out) in PLATES_STRAIGHT.items():
        y_in = INBOARD if y_out > 0 else -INBOARD
        path, span = _plate(stage, name, x0, x1, y_in, y_out)
        print(f"  {name:20s} straight  x[{span[0]:+.3f},{span[1]:+.3f}] "
              f"y[{span[2]:+.3f},{span[3]:+.3f}]")

    x0, x1 = -7.03, -6.39
    quad = [(x0, INBOARD), (x1, INBOARD), (x1, LANE_C_EDGE(x1)), (x0, LANE_C_EDGE(x0))]
    _trapezoid(stage, "PlowTransition_C", quad)
    print(f"  PlowTransition_C     trapezoid corners "
          + " ".join(f"({a:+.2f},{b:+.2f})" for a, b in quad))
    print(f"    outer edge follows the lane diagonal y = x + 7.637 "
          f"({LANE_C_EDGE(x0):+.3f} at x={x0}, {LANE_C_EDGE(x1):+.3f} at x={x1})")
    stage.GetRootLayer().Save()
    print(f"saved {SCENE}")
    print(f"driven surface now reaches |y| = {INBOARD}; the blade releases goods at ~0.26")


if __name__ == "__main__":
    main()