"""Bake catalogue scale and a clean origin into copies of the item meshes.

The originals are 2.0-2.8x smaller than the dims they are labelled with and sit ~0.6 m
above their own origin. Referencing them and fixing that with xform ops nests a scale
under a translate (which scales the translation) or a rigid body under a rigid body
(which hands the collider to the inner one and free-falls the outer). Baking the
transform into the points removes both traps: the result seats on z=0, is centred in
x/y, and loads through exactly the path run_pipeline.py already proves.
"""
import json, pathlib, shutil, sys
import numpy as np
from pxr import Usd, UsdGeom, Gf

SRC = pathlib.Path("/home/dasha/robozon-sorter/control_test/items")
DST = pathlib.Path("/home/dasha/robozon-sorter/control_test/items_flow")
DST.mkdir(exist_ok=True)
labels = json.loads((SRC / "labels.json").read_text())
shutil.copy(SRC / "labels.json", DST / "labels.json")
for t in ("textures",):
    if (SRC / t).exists() and not (DST / t).exists():
        try: (DST / t).symlink_to(SRC / t)
        except FileExistsError: pass

names = sys.argv[1:] or sorted(n for n in labels)
print(f"{'товар':20s} {'исходник, мм':22s} {'масштаб':>8s} {'после, мм':22s}")
for name in names:
    src = SRC / f"{name}.usd"
    if not src.exists():
        continue
    dst = DST / f"{name}.usd"
    shutil.copy(src, dst)
    st = Usd.Stage.Open(str(dst))
    root = st.GetDefaultPrim()
    xc = UsdGeom.XformCache(Usd.TimeCode.Default())
    meshes = [d for d in Usd.PrimRange(root) if d.IsA(UsdGeom.Mesh)]
    if not meshes:
        print(f"{name:20s} нет мешей, пропуск"); continue
    # world-space points under the item's own transforms
    allpts, per = [], []
    for m in meshes:
        P = np.asarray(UsdGeom.Mesh(m).GetPointsAttr().Get(), dtype=np.float64)
        M = np.array(xc.GetLocalToWorldTransform(m), dtype=np.float64)
        W = (np.c_[P, np.ones(len(P))] @ M)[:, :3]
        per.append((m, W)); allpts.append(W)
    A = np.concatenate(allpts, 0)
    lo, hi = A.min(0), A.max(0)
    ext = hi - lo
    s = max(labels[name]["dims_mm"]) / 1000.0 / float(max(ext))
    cx, cy = (lo[0] + hi[0]) / 2.0, (lo[1] + hi[1]) / 2.0
    BELT_CLEAR_MM = 430.0
    # Singulated pose: rest on the largest face and send the longest axis down the belt.
    # The line's belt is 450 mm between rails (measured: /World/_Rails at y +-0.22), so an
    # item presented across its long axis jams and the whole queue stops behind it.
    order = list(np.argsort(-ext))                     # largest -> X, middle -> Y, least -> Z
    for m, W in per:
        Q = (W - np.array([cx, cy, lo[2]]))[:, order] * s
        UsdGeom.Mesh(m).GetPointsAttr().Set([Gf.Vec3f(float(a), float(b), float(c)) for a, b, c in Q])
        e = UsdGeom.Mesh(m).GetExtentAttr()
        if e:
            e.Set([Gf.Vec3f(*[float(v) for v in Q.min(0)]), Gf.Vec3f(*[float(v) for v in Q.max(0)])])
        UsdGeom.Xformable(m).ClearXformOpOrder()       # transform is now in the points
    for d in Usd.PrimRange(root):
        if d.IsA(UsdGeom.Xformable):
            UsdGeom.Xformable(d).ClearXformOpOrder()
        if d.IsA(UsdGeom.Imageable):
            UsdGeom.Imageable(d).GetVisibilityAttr().Set(UsdGeom.Tokens.inherited)
    st.GetRootLayer().Save()
    chk = Usd.Stage.Open(str(dst))
    r = UsdGeom.BBoxCache(Usd.TimeCode.Default(),
                          [UsdGeom.Tokens.default_, UsdGeom.Tokens.render]
                          ).ComputeWorldBound(chk.GetDefaultPrim()).ComputeAlignedRange()
    out = [(r.GetMax()[k] - r.GetMin()[k]) * 1000 for k in range(3)]
    warn = "  ШИРЕ ПОЛОТНА" if out[1] > BELT_CLEAR_MM else ""
    print(f"{name:20s} {str([round(v*1000) for v in ext]):22s} {s:8.2f} "
          f"{str([round(v) for v in out]):22s} поперёк {out[1]:5.0f} мм{warn}")