#!/usr/bin/env python3
"""Place the two plow take-away lanes clear of the blade, and put a container at each end.

    python scripts/place_plow_lanes.py

The plow (x = -7.05, 600 mm arm, +-35 deg about Z) sweeps goods off the side of the main
run. A take-away lane therefore has to start OUTSIDE the main belt, not on it: the first
attempt put both near edges at y = +-0.05, which is inside the 450 mm belt, so the lanes
sat under the blade and fouled its swing.

Final layout, near edges at y = +-0.25 (just past the belt edge at +-0.225):

    lane C  +Y side, rotated 45 deg   near end (-7.05, +0.45)  runs toward (-X, +Y)
    lane B  -Y side, perpendicular    near end (-7.25, -0.45)  runs -Y
    container at the far end of each

Both sit at belt height (top z = 1.781) and start flush with ConveyorTrack_04's belt edge
at y = +-0.45, so they join the run the same way the D branch joins it upstream.

C is angled because a plow deflection carries goods sideways *and* downstream - they leave
the belt on a diagonal, and a 45 deg lane meets that trajectory instead of fighting it.
B stays square because the -Y throw is the shorter one.

Only these two tracks move. The plow, its hinge, its drive and the main run are untouched.
"""
from __future__ import annotations

import math
import shutil
import sys
from pathlib import Path

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

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

LANE_B = "/ConveyorTrack_01"           # -Y side, perpendicular
LANE_C = "/World/ConveyorTrack_01"     # +Y side, 45 deg

BELT_Z = 1.781
LANE_SCALE = Gf.Vec3d(1.0, 0.5, 1.0)   # same as every other track
# Height matters and was wrong the first time. The D branch (ConveyorTrack_03/Belt_01)
# sits at z 1.74..1.78, flush with the main belt, which is what makes it read as part of
# the conveyor. These tracks carried an authored -0.1 drop, putting them 100 mm low so they
# looked like separate furniture parked nearby. 0.0 puts their belt tops at 1.781 too.
LANE_DROP = 0.0

# Near ends must clear the ARM's swept envelope, not just the belt edge. The 600 mm arm
# pivots at (-7.05, 0) and reaches +-35 deg, so its tip traces out to
# 0.60*sin(35) = 0.344 m either side. Starting the lanes at +-0.45 leaves ~100 mm.
ARM_SWEEP_Y = 0.60 * math.sin(math.radians(35.0))       # 0.344 m
# A plow sweeps goods sideways while they are still ON the belt, so a take-away lane has to
# run ALONGSIDE it, its near edge touching the belt's side rail - not past the belt's end.
# The first placement put lane B at x -7.48..-7.03, entirely downstream of where the main
# belt stops (x = -7.00): goods would have had to leave the belt and cross a gap to reach
# it, which is why nothing ever arrived. Both lanes now sit inside the arm's working span
# (x -7.12..-6.52) with their near edges on the belt edge at y = +-0.45.
B_TRANSLATE = Gf.Vec3d(-6.80, -0.45, LANE_DROP)
B_YAW = -90.0                          # travel -Y
# A 45 deg lane meeting a straight belt edge does not touch at its centreline: the near
# corner runs ahead of it. Measured overlap at y=0.45 was 159 mm into the belt, so the lane
# is offset by that much and its nearest corner then lands on the edge instead of inside it.
# x=-6.55 puts the lane's near corner at the arm's tip (-6.52) and inside the belt span
# (-7.00..-6.00), i.e. on the junction itself. At -6.90 the corner sat behind the plow, so
# a +Y deflection had nowhere to land and goods rode on past.
C_TRANSLATE = Gf.Vec3d(-6.55, 0.45 + 0.159, LANE_DROP)
C_YAW = 135.0                          # travel (-X, +Y): the diagonal a plow throw makes


def yaw_quat(deg):
    r = math.radians(deg) / 2.0
    return Gf.Quatd(math.cos(r), Gf.Vec3d(0.0, 0.0, math.sin(r)))


def set_xform(layer, path, translate, yaw, scale=LANE_SCALE):
    spec = layer.GetPrimAtPath(path)
    if not spec:
        return False
    for name, value, vtype in (
            ("xformOp:translate", translate, Sdf.ValueTypeNames.Double3),
            ("xformOp:orient", yaw_quat(yaw), Sdf.ValueTypeNames.Quatd),
            ("xformOp:scale", scale, Sdf.ValueTypeNames.Double3)):
        attr = spec.attributes.get(name) or Sdf.AttributeSpec(spec, name, vtype)
        attr.default = value
    order = spec.attributes.get("xformOpOrder") or Sdf.AttributeSpec(
        spec, "xformOpOrder", Sdf.ValueTypeNames.TokenArray)
    order.default = ["xformOp:translate", "xformOp:orient", "xformOp:scale"]
    return True


# ---------------------------------------------------------------- containers
CONTAINERS = "/World/PlowContainers"
BIN_HALF = Gf.Vec3f(0.45, 0.40, 0.26)      # inner half-extents
BIN_FLOOR_Z = 1.16
BIN_WALL_TOP = 1.70                        # under the lane surface, so goods tip in
TH = 0.024


def _material(stage, path, rgb):
    mat = UsdShade.Material.Define(stage, path)
    sh = UsdShade.Shader.Define(stage, path + "/Shader")
    sh.CreateIdAttr("UsdPreviewSurface")
    sh.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(Gf.Vec3f(*rgb))
    sh.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.55)
    mat.CreateSurfaceOutput().ConnectToSource(sh.ConnectableAPI(), "surface")
    return mat


def _box(stage, path, centre, half, mat, collider=True):
    cube = UsdGeom.Cube.Define(stage, path)
    cube.GetSizeAttr().Set(2.0)                     # scale == half-extent
    xf = UsdGeom.Xformable(cube.GetPrim())
    xf.ClearXformOpOrder()
    xf.AddTranslateOp().Set(Gf.Vec3d(*centre))
    xf.AddScaleOp().Set(Gf.Vec3f(*half))
    UsdShade.MaterialBindingAPI.Apply(cube.GetPrim())
    UsdShade.MaterialBindingAPI(cube.GetPrim()).Bind(mat)
    if collider:
        UsdPhysics.CollisionAPI.Apply(cube.GetPrim())
    return cube.GetPrim()


def add_container(stage, tag, centre, mat):
    """open-top box; the far wall rises above belt height so a moving item cannot skim over"""
    cx, cy = centre
    hx, hy, _ = BIN_HALF
    wh = (BIN_WALL_TOP - BIN_FLOOR_Z) / 2
    wz = BIN_FLOOR_Z + wh
    back_top = BELT_Z + 0.22
    back_h = (back_top - BIN_FLOOR_Z) / 2
    base = f"{CONTAINERS}/{tag}"
    _box(stage, f"{base}_Floor", (cx, cy, BIN_FLOOR_Z), (hx, hy, TH), mat)
    for name, c, h in (
            ("W0", (cx, cy + hy, BIN_FLOOR_Z + back_h), (hx, TH, back_h)),
            ("W1", (cx, cy - hy, wz), (hx, TH, wh)),
            ("W2", (cx - hx, cy, wz), (TH, hy, wh)),
            ("W3", (cx + hx, cy, wz), (TH, hy, wh))):
        _box(stage, f"{base}_{name}", c, h, mat)
    for i, (lx, ly) in enumerate([(cx - hx + 0.06, cy - hy + 0.06), (cx + hx - 0.06, cy - hy + 0.06),
                                  (cx - hx + 0.06, cy + hy - 0.06), (cx + hx - 0.06, cy + hy - 0.06)]):
        _box(stage, f"{base}_Leg{i}", (lx, ly, BIN_FLOOR_Z / 2),
             (0.024, 0.024, BIN_FLOOR_Z / 2), mat, collider=False)


# ---------------------------------------------------------------- transition plate
def add_transition(stage, lane_path, edge_y, sign, mat, edge_x0, edge_x1):
    """Deck the WHOLE discharge corner, not just the touching triangle.

    An angled lane leaves gaps on BOTH sides of its end face: one between its near corner
    and the belt edge, another beyond its far corner. Goods do not cross at a single point -
    the plow can put them anywhere across the discharge width - so the plate has to span the
    entire region between the belt edge and the lane's end face, from one side of the zone
    to the other. A plate covering only the first triangle still drops anything pushed wide.

    Built as the convex hull of the belt-edge segment and both end-face corners, extruded
    down 20 mm, coplanar with both belt surfaces.
    """
    cache = UsdGeom.BBoxCache(0, ["default", "render"], useExtentsHint=True)
    xf = UsdGeom.XformCache()
    belt = stage.GetPrimAtPath(f"{lane_path}/Belt")
    M = xf.GetLocalToWorldTransform(belt)
    r = cache.ComputeWorldBound(belt).ComputeAlignedRange()
    top = r.GetMax()[2]

    travel = M.TransformDir(Gf.Vec3d(1, 0, 0)).GetNormalized()
    across = Gf.Vec3d(-travel[1], travel[0], 0.0)
    half_w = 0.225
    centre = Gf.Vec3d((r.GetMin()[0] + r.GetMax()[0]) / 2,
                      (r.GetMin()[1] + r.GetMax()[1]) / 2, 0.0)
    near = centre - travel * 1.0
    c0 = near + across * half_w
    c1 = near - across * half_w

    # everything the plate must reach: the belt edge across the discharge zone, and both
    # corners of the lane's end face
    xs = [edge_x0, edge_x1, float(c0[0]), float(c1[0])]
    ys = [edge_y, edge_y, float(c0[1]), float(c1[1])]
    far_y = max(ys) if sign > 0 else min(ys)
    quad = [(min(xs), edge_y), (max(xs), edge_y), (max(xs), far_y), (min(xs), far_y)]

    path = f"/World/PlowTransition_{'C' if sign > 0 else 'B'}"
    if stage.GetPrimAtPath(path).IsValid():
        stage.RemovePrim(path)
    if abs(far_y - edge_y) < 0.002:
        # a square lane meets the edge flush along its whole face - no gap to deck, and a
        # zero-area mesh would be a degenerate collider
        return None, []
    mesh = UsdGeom.Mesh.Define(stage, path)
    pts = [Gf.Vec3f(x, y, top) for x, y in quad] + [Gf.Vec3f(x, y, top - 0.02) for x, y in quad]
    mesh.GetPointsAttr().Set(pts)
    faces = [(0, 3, 2, 1), (4, 5, 6, 7), (0, 1, 5, 4), (1, 2, 6, 5), (2, 3, 7, 6), (3, 0, 4, 7)]
    mesh.GetFaceVertexCountsAttr().Set([4] * len(faces))
    mesh.GetFaceVertexIndicesAttr().Set([i for f in faces for i in f])
    mesh.GetSubdivisionSchemeAttr().Set("none")
    mesh.GetExtentAttr().Set([Gf.Vec3f(min(xs), min(edge_y, far_y), top - 0.02),
                              Gf.Vec3f(max(xs), max(edge_y, far_y), top)])
    UsdShade.MaterialBindingAPI.Apply(mesh.GetPrim())
    UsdShade.MaterialBindingAPI(mesh.GetPrim()).Bind(mat)
    UsdPhysics.CollisionAPI.Apply(mesh.GetPrim())
    return path, [(round(x, 3), round(y, 3)) for x, y in quad]


def add_corner_deck(stage, lane_path, edge_y, sign, mat, edge_far_x):
    """Close the right angle a SQUARE lane leaves against a wider belt.

    Lane B meets the belt flush along its own face, but the belt is wider than the lane:
    the run reaches x=-6.00 while the lane stops at x=-6.58. Anything the plow pushes
    sideways in that leftover span has open air under it. A triangular fillet spanning the
    belt edge out to the run's end and back down the lane's side turns that right angle
    into a chute, so goods slide into the lane instead of dropping through.
    """
    cache = UsdGeom.BBoxCache(0, ["default", "render"], useExtentsHint=True)
    belt = stage.GetPrimAtPath(f"{lane_path}/Belt")
    r = cache.ComputeWorldBound(belt).ComputeAlignedRange()
    top = r.GetMax()[2]
    lane_far_x = r.GetMax()[0]              # the lane's edge nearest the open span
    span = abs(edge_far_x - lane_far_x)
    if span < 0.02:
        return None, []
    corner = (lane_far_x, edge_y)
    along_belt = (edge_far_x, edge_y)
    down_lane = (lane_far_x, edge_y + sign * span)

    path = f"/World/PlowCornerDeck_{'C' if sign > 0 else 'B'}"
    if stage.GetPrimAtPath(path).IsValid():
        stage.RemovePrim(path)
    tri = [corner, along_belt, down_lane]
    mesh = UsdGeom.Mesh.Define(stage, path)
    pts = [Gf.Vec3f(x, y, top) for x, y in tri] + [Gf.Vec3f(x, y, top - 0.02) for x, y in tri]
    mesh.GetPointsAttr().Set(pts)
    faces = [(0, 2, 1), (3, 4, 5), (0, 1, 4, 3), (1, 2, 5, 4), (2, 0, 3, 5)]
    mesh.GetFaceVertexCountsAttr().Set([len(f) for f in faces])
    mesh.GetFaceVertexIndicesAttr().Set([i for f in faces for i in f])
    mesh.GetSubdivisionSchemeAttr().Set("none")
    xs = [p[0] for p in tri]; ys = [p[1] for p in tri]
    mesh.GetExtentAttr().Set([Gf.Vec3f(min(xs), min(ys), top - 0.02),
                              Gf.Vec3f(max(xs), max(ys), top)])
    UsdShade.MaterialBindingAPI.Apply(mesh.GetPrim())
    UsdShade.MaterialBindingAPI(mesh.GetPrim()).Bind(mat)
    UsdPhysics.CollisionAPI.Apply(mesh.GetPrim())
    return path, [(round(x, 3), round(y, 3)) for x, y in tri]


def main():
    if not CELL.exists():
        print(f"{CELL} not found")
        return 1
    backup = CELL.with_suffix(".usd.prelanes")
    if not backup.exists():
        shutil.copy(CELL, backup)
        print(f"backup -> {backup.name}")

    layer = Sdf.Layer.FindOrOpen(str(CELL))
    set_xform(layer, LANE_B, B_TRANSLATE, B_YAW)
    set_xform(layer, LANE_C, C_TRANSLATE, C_YAW)
    layer.Save()
    print(f"arm sweeps to y=+-{ARM_SWEEP_Y:.3f}; lanes start at +-0.45")
    print(f"lane B (-Y, square)  near end {tuple(B_TRANSLATE)[:2]}")
    print(f"lane C (+Y, 45 deg)  near end {tuple(C_TRANSLATE)[:2]}")

    stage = Usd.Stage.Open(str(CELL))
    cache = UsdGeom.BBoxCache(0, ["default", "render"], useExtentsHint=True)
    xf = UsdGeom.XformCache()

    if stage.GetPrimAtPath(CONTAINERS).IsValid():
        stage.RemovePrim(CONTAINERS)
    UsdGeom.Xform.Define(stage, CONTAINERS)
    m_b = _material(stage, f"{CONTAINERS}/M_B", (0.85, 0.22, 0.20))
    m_c = _material(stage, f"{CONTAINERS}/M_C", (0.25, 0.72, 0.32))

    ends = {}
    for tag, path in (("B", LANE_B), ("C", LANE_C)):
        belt = stage.GetPrimAtPath(f"{path}/Belt")
        r = cache.ComputeWorldBound(belt).ComputeAlignedRange()
        mn, mx = r.GetMin(), r.GetMax()
        d = xf.GetLocalToWorldTransform(belt).TransformDir(Gf.Vec3d(1, 0, 0)).GetNormalized()
        centre = Gf.Vec3d((mn[0] + mx[0]) / 2, (mn[1] + mx[1]) / 2, 0)
        # container just past the discharge end, along the lane's own travel direction
        far = centre + d * ((max(mx[0] - mn[0], mx[1] - mn[1]) / 2) + BIN_HALF[1] + 0.10)
        ends[tag] = (float(far[0]), float(far[1]))
        print(f"  lane {tag}: x[{mn[0]:6.2f}..{mx[0]:6.2f}] y[{mn[1]:6.2f}..{mx[1]:6.2f}] "
              f"top={mx[2]:.3f} travel({d[0]:+.2f},{d[1]:+.2f})")

    add_container(stage, "B", ends["B"], m_b)
    add_container(stage, "C", ends["C"], m_c)

    m_t = _material(stage, f"{CONTAINERS}/M_transition", (0.30, 0.31, 0.34))
    for old in ("/World/PlowTransition_C", "/World/PlowTransition_B"):
        if stage.GetPrimAtPath(old).IsValid():
            stage.RemovePrim(old)
    # deck both discharge corners across the full width of the junction (belt x -7.00..-6.00)
    for lane, edge, sign in ((LANE_C, 0.45, +1), (LANE_B, -0.45, -1)):
        path, corners = add_transition(stage, lane, edge, sign, m_t, -7.00, -6.00)
        print(f"  transition {path or 'not needed (lane meets flush)'}: {corners}")
        # a flush lane still leaves the right angle where the belt runs on past it
        cpath, ctri = add_corner_deck(stage, lane, edge, sign, m_t, -6.00)
        print(f"  corner deck {cpath or 'not needed'}: {ctri}")
    stage.GetRootLayer().Save()

    print("\ncontainers:")
    for tag in ("B", "C"):
        r = cache.ComputeWorldBound(
            stage.GetPrimAtPath(f"{CONTAINERS}/{tag}_Floor")).ComputeAlignedRange()
        mn, mx = r.GetMin(), r.GetMax()
        print(f"  {tag}: x[{mn[0]:6.2f}..{mx[0]:6.2f}] y[{mn[1]:6.2f}..{mx[1]:6.2f}] "
              f"floor z={mx[2]:.2f}")

    arm = cache.ComputeWorldBound(
        stage.GetPrimAtPath("/World/Diverters/DiverterEnd/Arm")).ComputeAlignedRange()
    print(f"\nplow arm x[{arm.GetMin()[0]:.2f}..{arm.GetMax()[0]:.2f}] "
          f"y[{arm.GetMin()[1]:.2f}..{arm.GetMax()[1]:.2f}]; swept envelope +-{ARM_SWEEP_Y:.2f}")
    return 0


if __name__ == "__main__":
    sys.exit(main())