#!/usr/bin/env python3
"""Bring the plow arm to a 600 mm sweep width.

    python scripts/narrow_plow.py [--width 0.60]

The authored arm is 730 mm along its own long axis. That is wider than the 450 mm belt by
enough that it overhangs both rails: it clips goods it should have passed and shoulders
others off the lane instead of steering them. 600 mm still spans the belt with margin while
leaving the lane edges clear.

Only the scale on `DiverterEnd/Arm/Geom` changes. The hinge, its drive, the limits and the
arm's rigid body are untouched, so the kinematics are exactly as authored - the arm is
simply shorter.

Which local axis carries the length is not obvious: Geom is rotated -90 deg about Z and its
parent another 180 deg, so the mesh's own X and Y do not map to the world axes you would
guess. The script measures instead of assuming.
"""
from __future__ import annotations

import argparse
import shutil
import sys
from pathlib import Path

import numpy as np
from pxr import Gf, Usd, UsdGeom

ROOT = Path(__file__).resolve().parent.parent
CELL = ROOT / "scene" / "plow_cell.usd"
ARM_GEOM = "/World/Diverters/DiverterEnd/Arm/Geom"
ARM = "/World/Diverters/DiverterEnd/Arm"


def arm_length(stage):
    """longest principal extent of the arm's mesh points, in world metres"""
    cache = UsdGeom.XformCache()
    pts = []
    for prim in Usd.PrimRange(stage.GetPrimAtPath(ARM)):
        mesh = UsdGeom.Mesh(prim)
        if not mesh:
            continue
        p = mesh.GetPointsAttr().Get()
        if not p:
            continue
        M = cache.GetLocalToWorldTransform(prim)
        pts.append(np.array([M.Transform(Gf.Vec3d(*q)) for q in p]))
    if not pts:
        return None
    P = np.vstack(pts)
    Q = P - P.mean(0)
    _, _, vt = np.linalg.svd(Q, full_matrices=False)
    return float(np.ptp(Q @ vt[0]))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--width", type=float, default=0.60, help="target sweep width, metres")
    args = ap.parse_args()

    if not CELL.exists():
        print(f"{CELL} not found")
        return 1

    stage = Usd.Stage.Open(str(CELL))
    geom = stage.GetPrimAtPath(ARM_GEOM)
    if not geom.IsValid():
        print(f"{ARM_GEOM} missing - is this plow_cell.usd?")
        return 1

    before = arm_length(stage)
    if not before:
        print("arm carries no mesh points - is assets/plow/ populated?")
        return 1
    print(f"arm length now {before*1000:.0f} mm, target {args.width*1000:.0f} mm")

    xf = UsdGeom.Xformable(geom)
    scale_op = None
    for op in xf.GetOrderedXformOps():
        if op.GetOpType() == UsdGeom.XformOp.TypeScale:
            scale_op = op
    if scale_op is None:
        scale_op = xf.AddScaleOp()
        scale_op.Set(Gf.Vec3f(1, 1, 1))
    base = Gf.Vec3f(scale_op.Get() or Gf.Vec3f(1, 1, 1))

    # find which local axis the length rides on, by testing rather than reasoning about
    # the two stacked rotations
    factor = args.width / before
    best = None
    for axis in (0, 1, 2):
        trial = Gf.Vec3f(base)
        trial[axis] = base[axis] * factor
        scale_op.Set(trial)
        got = arm_length(stage)
        print(f"  scale on local {'XYZ'[axis]} -> {got*1000:.0f} mm")
        if best is None or abs(got - args.width) < abs(best[1] - args.width):
            best = (axis, got, trial)
    axis, got, trial = best
    scale_op.Set(trial)

    if abs(got - args.width) > 0.005:
        print(f"closest achievable was {got*1000:.0f} mm on local {'XYZ'[axis]} - "
              "the arm's length may not lie on a single local axis")
        return 1

    backup = CELL.with_suffix(".usd.prewidth")
    if not backup.exists():
        shutil.copy(CELL, backup)
        print(f"backup -> {backup.name}")
    stage.GetRootLayer().Save()

    check = Usd.Stage.Open(str(CELL))
    final = arm_length(check)
    cache = UsdGeom.BBoxCache(0, ["default", "render"], useExtentsHint=True)
    r = cache.ComputeWorldBound(check.GetPrimAtPath(ARM)).ComputeAlignedRange()
    mn, mx = r.GetMin(), r.GetMax()
    print(f"\nsaved. arm is now {final*1000:.0f} mm "
          f"(scale {tuple(round(v,4) for v in trial)} on local {'XYZ'[axis]})")
    print(f"  world AABB x[{mn[0]:.3f}..{mx[0]:.3f}] y[{mn[1]:.3f}..{mx[1]:.3f}] "
          f"z[{mn[2]:.3f}..{mx[2]:.3f}]")
    return 0


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