LOPs - PY - advanced instancer

from pxr import Usd, UsdGeom, Gf, Vt, Sdf
import numpy as np
import hou

def resolve_proto_indices_numpy(proto_names, prototype_lookup, fallback_index=0):
    if not proto_names:
        return []
    proto_names_np = np.array(proto_names, dtype=object)
    name_to_index = {k: prototype_lookup.get(k, fallback_index) for k in proto_names_np}
    indices_np = np.vectorize(name_to_index.get)(proto_names_np)
    return list(indices_np)

# --- Houdini Node Context ---
node = hou.pwd()
stage = node.editableStage()

instancer_path = node.evalParm("instancer_path")
points_path = node.evalParm("points_path")
collection_path = node.evalParm("collection_path")
collection_name = node.evalParm("collection_name")

# --- Step 1: Construct full collection path and get collection prim ---
collection_base = collection_path
collection_prim_path = Sdf.Path(collection_base)
collection_full_name = f"collection:{collection_name}"
collection_path = f"{collection_base}.{collection_full_name}"

collection_prim = stage.GetPrimAtPath(collection_prim_path)
if not collection_prim or not collection_prim.IsValid():
    raise RuntimeError(f"Prim not found at path: {collection_prim_path}")

includes_rel = collection_prim.GetRelationship(f"{collection_full_name}:includes")
if not includes_rel or not includes_rel.IsValid():
    raise RuntimeError(f"'includes' relationship not found on collection: {collection_path}")

prototype_paths = [target.pathString for target in includes_rel.GetTargets()]
if not prototype_paths:
    raise RuntimeError(f"No prototype targets found in collection: {collection_path}")

# --- Step 2: Build prototype name → index lookup ---
prototype_lookup = {
    path.split("/")[-1]: idx
    for idx, path in enumerate(prototype_paths)
}

# --- Step 3: Get point attributes from source prim ---
points_prim = stage.GetPrimAtPath(points_path)
if not points_prim or not points_prim.IsValid():
    raise RuntimeError(f"Invalid points prim path: {points_path}")

points_geom = UsdGeom.Points(points_prim)

# Positions
positions = points_geom.GetPointsAttr().Get() or []

# Orientations
orient_attr = points_prim.GetAttribute("orient")
if orient_attr and orient_attr.HasValue():
    raw_orients = orient_attr.Get()
    orientations = [
        Gf.Quath(q.real, *q.imaginary) if hasattr(q, "real") else
        Gf.Quath(q.GetReal(), *q.GetImaginary())
        for q in raw_orients
    ]
else:
    orientations = [Gf.Quath(1.0, 0.0, 0.0, 0.0) for _ in positions]

# Scales
scale_attr = points_prim.GetAttribute("primvars:scale")
if scale_attr and scale_attr.HasValue():
    raw_scales = scale_attr.Get()
    scales = [Gf.Vec3f(*s) if hasattr(s, "__iter__") else s for s in raw_scales]
else:
    scales = [Gf.Vec3f(1.0, 1.0, 1.0) for _ in positions]

# Proto Names → Indices
proto_name_attr = points_prim.GetAttribute("primvars:protoName")
proto_names = proto_name_attr.Get() if proto_name_attr and proto_name_attr.HasValue() else []
proto_indices = resolve_proto_indices_numpy(proto_names, prototype_lookup)

# *** Fix: convert numpy.int64 to native int ***
proto_indices = [int(i) for i in proto_indices]

# --- Step 4: Create the PointInstancer ---
instancer_prim = stage.DefinePrim(instancer_path, "PointInstancer")
instancer = UsdGeom.PointInstancer(instancer_prim)

# Add a relationship to the source points prim
#source_points_rel = instancer_prim.CreateRelationship("sourcePoints")
#source_points_rel.SetTargets([points_prim.GetPath()])

instancer.GetPrototypesRel().SetTargets(prototype_paths)
instancer.GetProtoIndicesAttr().Set(Vt.IntArray(proto_indices))
instancer.GetPositionsAttr().Set(positions)
instancer.GetOrientationsAttr().Set(orientations)
instancer.GetScalesAttr().Set(scales)

# --- Step 5: Make all prototypes invisible ---
#for proto_path in prototype_paths:
#    proto_prim = stage.GetPrimAtPath(proto_path)
#    if proto_prim and proto_prim.IsValid():
#        UsdGeom.Imageable(proto_prim).MakeInvisible()