"""STAGE 1 (inside Isaac, no torch): render rectified L/R pairs + an EXACT object mask.

Two fixes over capture_cfg.py:
  * the reference used to be composed onto the same prim whose xformOpOrder we then cleared,
    which destroyed the mesh's own placement and dropped every item 0.6 m under the belt.
    The reference now lives on a child, so only our holder carries the placement.
  * ground truth is the mesh's real bbox in the scene, not the product catalogue - the two
    disagree by 2.0-2.8x per item, so catalogue dims cannot score a size prediction.

The mask is analytic (mesh points projected through the camera), the same trick the earlier
ROI pipeline used for its GT channel - no segmentation error contaminates a geometry study.
"""
import json, os, sys
REPO = "/home/dasha/robozon-sorter"
for e in (REPO, f"{REPO}/control_test"):
    if e not in sys.path:
        sys.path.insert(0, e)
for _m in [k for k in list(sys.modules) if k.startswith(("cam_configs", "classify", "cell"))]:
    del sys.modules[_m]
import importlib; importlib.invalidate_caches()

import asyncio, numpy as np, cv2
import omni.usd, omni.timeline
import omni.kit.viewport.utility as vp
import isaacsim.core.experimental.utils.app as app_utils
from pxr import UsdPhysics, Gf, Usd, UsdGeom, UsdLux

import cam_configs as CC
import cell
import classify as CL

try:                       # --file runs isolated, so an injected arg lands in
    CFG = cfg              # LOCALS, not globals() - the bare name catches both
except NameError:
    CFG = CC.DEFAULT
ITEMS = globals().get("items", ["bag", "backpack", "lunchbox", "helmet", "pillow",
                                "detergent", "bucket", "box_400x400x300", "box_300x200x200"])
OUT = f"{REPO}/control_test/captures/{CFG}"
os.makedirs(OUT, exist_ok=True)

stage = omni.usd.get_context().get_stage()
tl = omni.timeline.get_timeline_interface()
if tl.is_playing():
    tl.stop(); await app_utils.update_app_async(steps=10)

calib = CC.apply_config(stage, CFG)
print(f"config {CFG}: {len(calib)} камер | " + ", ".join(
    f"{n.replace('_Left','')} h={c['height_mm']:.0f} d={c['standoff_mm']:.0f}"
    for n, c in calib.items() if n.endswith("_Left")))

# without this the cell has no floor and no dome: everything off the belt renders as
# void, which both looks wrong over WebRTC and starves the side views of bounce light
_gl = cell.add_ground_and_light(stage)
print(f"пол {_gl['ground']}, купол {_gl['light']}")

for nm, pos in (("K0", (-0.75, 1.6, 2.6)), ("K1", (-0.75, -1.6, 2.6)),
                ("K2", (0.6, 0.0, 2.6)), ("K3", (-2.1, 0.0, 2.6))):
    p = f"/World/_CapLight_{nm}"
    if not stage.GetPrimAtPath(p).IsValid():
        sl = UsdLux.SphereLight.Define(stage, p)
        sl.CreateRadiusAttr().Set(0.25); sl.CreateIntensityAttr().Set(90000.0)
        UsdGeom.Xformable(sl.GetPrim()).AddTranslateOp().Set(Gf.Vec3d(*pos))
dome = stage.GetPrimAtPath("/Environment/_BrightFill")
if dome.IsValid():
    dome.GetAttribute("inputs:intensity").Set(3500.0)

lib = {r["name"]: r for r in CL.load_library(f"{REPO}/control_test/items") if "error" not in r}
ROOT = "/World/CapItems2"
UsdGeom.Xform.Define(stage, ROOT)
BB = lambda: UsdGeom.BBoxCache(Usd.TimeCode.Default(),
                               [UsdGeom.Tokens.default_, UsdGeom.Tokens.render])

async def spawn(name):
    """holder Xform carries the placement; the reference sits on a child so its own
    transform survives."""
    path = f"{ROOT}/{name}"
    if stage.GetPrimAtPath(path).IsValid():
        stage.RemovePrim(path)
    holder = UsdGeom.Xform.Define(stage, path).GetPrim()
    inner = UsdGeom.Xform.Define(stage, f"{path}/mesh").GetPrim()
    inner.GetReferences().AddReference(lib[name]["path"])
    r = None
    for _ in range(12):                      # a reference does not compose within the tick
        await app_utils.update_app_async(steps=2)
        r = BB().ComputeWorldBound(inner).ComputeAlignedRange()
        if not r.IsEmpty():
            break
    if r is None or r.IsEmpty():
        raise RuntimeError(f"{name}: пустой bbox - ссылка не разрешилась")
    mn, mx = r.GetMin(), r.GetMax()
    UsdGeom.Xformable(holder).AddTranslateOp(precision=UsdGeom.XformOp.PrecisionDouble).Set(
        Gf.Vec3d(CC.TARGET[0] - (mn[0] + mx[0]) / 2.0,
                 CC.TARGET[1] - (mn[1] + mx[1]) / 2.0,
                 CC.TARGET[2] - mn[2] + 0.001))
    # the exported item layers author visibility=invisible on their own root, and
    # MakeVisible on an ancestor does NOT clear a descendant's authored value - that is
    # why every earlier capture showed bare belt
    for d in Usd.PrimRange(holder):
        if d.IsA(UsdGeom.Imageable):
            UsdGeom.Imageable(d).GetVisibilityAttr().Set(UsdGeom.Tokens.inherited)
    UsdGeom.Imageable(holder).MakeVisible()
    r2 = BB().ComputeWorldBound(inner).ComputeAlignedRange()
    ext = [(r2.GetMax()[i] - r2.GetMin()[i]) * 1000.0 for i in range(3)]
    return holder, ext, [r2.GetMin()[i] for i in range(3)], [r2.GetMax()[i] for i in range(3)]

def world_geom(prim):
    """every mesh of the item in world space, as vertices + triangles. Vertices alone are
    not enough: a box has eight of them, so a point-splat mask covers ~60 px and the ROI
    collapses. Filling the projected triangles gives the true silhouette."""
    xc = UsdGeom.XformCache(Usd.TimeCode.Default()); V = []; T = []; base = 0
    for d in Usd.PrimRange(prim):
        if not d.IsA(UsdGeom.Mesh):
            continue
        m = UsdGeom.Mesh(d)
        pts = m.GetPointsAttr().Get()
        if not pts:
            continue
        M = np.array(xc.GetLocalToWorldTransform(d), dtype=np.float64)
        P = np.asarray(pts, dtype=np.float64)
        V.append((np.c_[P, np.ones(len(P))] @ M)[:, :3])
        cnt = m.GetFaceVertexCountsAttr().Get() or []
        idx = m.GetFaceVertexIndicesAttr().Get() or []
        o = 0
        for c in cnt:                                    # fan-triangulate each polygon
            for k in range(1, c - 1):
                T.append((base + idx[o], base + idx[o + k], base + idx[o + k + 1]))
            o += c
        base += len(P)
    if not V:
        return np.zeros((0, 3)), np.zeros((0, 3), int)
    return np.concatenate(V, 0), np.asarray(T, dtype=np.int64).reshape(-1, 3)


def gt_mask(V, T, cam):
    W, H = cam["width"], cam["height"]
    Minv = np.linalg.inv(np.array(cam["M"]))
    c = (np.c_[V, np.ones(len(V))] @ Minv)[:, :3]
    z = -c[:, 2]
    u = c[:, 0] / np.maximum(z, 1e-9) * cam["fx"] + cam["cx"]
    v = -c[:, 1] / np.maximum(z, 1e-9) * cam["fy"] + cam["cy"]
    uv = np.c_[u, v]
    m = np.zeros((H, W), np.uint8)
    if len(T):
        good = (z[T] > 1e-3).all(1)
        tri = uv[T[good]].astype(np.int32)
        tri = np.clip(tri, [-4 * W, -4 * H], [4 * W, 4 * H])
        cv2.fillPoly(m, list(tri), 1)
    else:
        ok = (z > 1e-3) & (u >= 0) & (u < W) & (v >= 0) & (v < H)
        m[v[ok].astype(int), u[ok].astype(int)] = 1
    m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, np.ones((5, 5), np.uint8))
    n, lab, st, _ = cv2.connectedComponentsWithStats(m)
    if n > 1:
        m = (lab == (1 + np.argmax(st[1:, cv2.CC_STAT_AREA]))).astype(np.uint8)
    return m.astype(bool)


w = vp.get_active_viewport(); orig_cam = w.camera_path
manifest = {"config": CFG, "target": [float(v) for v in CC.TARGET],
            "standoff_m": CC.STANDOFF, "calib": calib, "items": {}}

for name in ITEMS:
    if name not in lib:
        print(f"  пропуск {name} (нет в библиотеке)"); continue
    prim, ext, bmin, bmax = await spawn(name)
    await app_utils.update_app_async(steps=20)
    VV, TT = world_geom(prim)
    files, masks, cover, seen = {}, {}, {}, {}
    for cam_name, cam in calib.items():
        w.camera_path = f"/RigRS/{cam_name}" if stage.GetPrimAtPath(f"/RigRS/{cam_name}").IsValid() \
                        else CC._cam(stage, cam_name).GetPath()
        await app_utils.update_app_async(steps=22); await asyncio.sleep(0)
        f = f"{OUT}/{name}__{cam_name}.png"
        vp.capture_viewport_to_file(w, file_path=f)
        await app_utils.update_app_async(steps=12); await asyncio.sleep(0)
        files[cam_name] = f
        mk = gt_mask(VV, TT, cam)
        _img = cv2.imread(f)
        if _img is not None and _img.shape[:2] == mk.shape:
            _g = cv2.cvtColor(_img, cv2.COLOR_BGR2GRAY).astype(float)
            _ring = cv2.dilate(mk.astype(np.uint8), np.ones((41, 41), np.uint8)).astype(bool) & ~mk
            seen[cam_name] = round(float(abs(_g[mk].mean() - _g[_ring].mean())), 1)
        np.savez_compressed(f"{OUT}/{name}__{cam_name}_mask.npz", m=mk)
        masks[cam_name] = f"{OUT}/{name}__{cam_name}_mask.npz"
        cover[cam_name] = int(mk.sum())
    manifest["items"][name] = dict(
        files=files, masks=masks, mask_px=cover, nverts=int(len(VV)), ntris=int(len(TT)),
        contrast=seen, gt_catalogue=lib[name]["dims_mm"], cls=lib[name]["cls"],
        gt_scene_mm=[round(e, 1) for e in ext], bmin=bmin, bmax=bmax)
    print(f"  {name}: в сцене {[round(e) for e in sorted(ext, reverse=True)]} мм "
          f"маска {min(cover.values())}-{max(cover.values())} px, контраст {min(seen.values()):.0f}-{max(seen.values()):.0f}"
          + ("  <-- НЕ ВИДЕН" if min(seen.values()) < 3 else ""))
    # Спрятать МАЛО: невидимость не убирает коллайдер, и снятый товар остаётся твёрдой
    # стеной ровно в точке осмотра, посреди рабочей линии. После двух прогонов там стояло
    # 18 невидимых предметов, и поток вставал на них, не доезжая до плуга.
    UsdGeom.Imageable(prim).MakeInvisible()
    for _d in Usd.PrimRange(prim):
        _a = _d.GetAttribute("physics:collisionEnabled")
        if _a and _a.IsValid():
            _a.Set(False)
        elif _d.HasAPI(UsdPhysics.CollisionAPI):
            UsdPhysics.CollisionAPI(_d).CreateCollisionEnabledAttr().Set(False)

w.camera_path = orig_cam
await app_utils.update_app_async(steps=10)
json.dump(manifest, open(f"{OUT}/manifest.json", "w"), indent=1)
print(f"\nmanifest -> {OUT}/manifest.json")