@needle-tools/materialx 1.0.1-next.b9638d9 → 1.0.1-next.c1bbe8d

package/README.md

@@ -1,5 +1,10 @@
 # Needle MaterialX
 
+Load and display MaterialX materials in Needle Engine
+
+## Installation
+`npm i @needle-tools/materialx`  
+
 ## How to use
 
 To use with Needle Engine simply import the module

package/codegen/register_types.ts

@@ -2,5 +2,7 @@
 import { TypeStore } from "@needle-tools/engine"
 
 // Import types
+import { MaterialXUniformUpdate } from "../src/loader/loader.needle.js";
 
 // Register types
+TypeStore.add("MaterialXUniformUpdate", MaterialXUniformUpdate);

package/index.ts

@@ -2,4 +2,4 @@ import { registerNeedleLoader } from "./src/loader/loader.needle.js";
 
 registerNeedleLoader();
 
-export { initializeMaterialX, getMaterialXEnvironment } from "./src/index.js";
+export { ready, getMaterialXEnvironment } from "./src/index.js";

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@needle-tools/materialx",
-  "version": "1.0.1-next.b9638d9",
+  "version": "1.0.1-next.c1bbe8d",
   "type": "module",
   "main": "index.ts",
   "exports": {
@@ -8,7 +8,11 @@
       "import": "./index.ts",
       "require": "./index.js"
     },
-    "./package.json": "./package.json"
+    "./package.json": "./package.json",
+    "./codegen/register_types.ts": {
+      "import": "./codegen/register_types.ts",
+      "require": "./codegen/register_types.js"
+    }
   },
   "peerDependencies": {
     "@needle-tools/engine": "4.x",
@@ -23,5 +27,16 @@
   "publishConfig": {
     "access": "public",
     "registry": "https://registry.npmjs.org/"
-  }
+  },
+  "keywords": [
+    "needle",
+    "materialx",
+    "material",
+    "shader",
+    "threejs",
+    "three.js",
+    "webgl",
+    "mtlx",
+    "rendering"
+  ]
 }

package/src/helper.js

@@ -3,7 +3,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-import { getParam } from '@needle-tools/engine';
+import { getParam, getWorldDirection } from '@needle-tools/engine';
 import * as THREE from 'three';
 
 const debug = getParam("debugmaterialx");
@@ -88,7 +88,7 @@ function fromMatrix(matrix, dimension)
  * @param {mx.Uniform.value} value
  * @param {mx.Uniform.name} name
  * @param {mx.Uniforms} uniforms
- * @param {THREE.textureLoader} textureLoader
+ * @param {THREE.TextureLoader} textureLoader
  * @param {string} searchPath
  * @param {boolean} flipY
  */
@@ -291,19 +291,21 @@ export function findLights(doc)
     return lights;
 }
 
+let lightTypesBound = {};
+
 /**
  * Register lights in shader generation context
  * @param {Object} mx MaterialX Module
- * @param {Array.<mx.Node>} lights Light nodes
  * @param {mx.GenContext} genContext Shader generation context
  * @returns {Array.<mx.Node>}
  */
-export async function registerLights(mx, lights, genContext)
+export async function registerLights(mx, genContext)
 {
+    lightTypesBound = {};
+    const maxLightCount = genContext.getOptions().hwMaxActiveLightSources;
+
     mx.HwShaderGenerator.unbindLightShaders(genContext);
 
-    const lightTypesBound = {};
-    const lightData = [];
     let lightId = 1;
 
     // All light types so that we have NodeDefs for them
@@ -329,7 +331,11 @@ export async function registerLights(mx, lights, genContext)
     document.setDataLibrary(stdlib);
     document.importLibrary(lightRigDoc);
     const defaultLights = findLights(document);
-    if (debug) console.log("Default lights in MaterialX document", defaultLights);
+    // if (debug) 
+        console.log("Default lights in MaterialX document", defaultLights);
+
+    // Loading a document seems to reset this option for some reason, so we set it again
+    genContext.getOptions().hwMaxActiveLightSources = maxLightCount;
 
     // Register types only – we get these from the default light rig XML above
     // This is needed to ensure that the light shaders are bound for each light type
@@ -351,52 +357,34 @@ export async function registerLights(mx, lights, genContext)
     }
 
     if (debug) console.log("Light types bound in MaterialX context", lightTypesBound);
+}
 
-    // MaterialX light nodes
-    for (let light of lights)
-    {
-        // Skip if light does not have a node definition
-        if (!("getNodeDef" in light)) continue;
-
-        let nodeDef = light.getNodeDef();
-        let nodeName = nodeDef.getName();
-        if (!lightTypesBound[nodeName])
-        {
-            if (debug) console.log("bind light shader for node", { nodeName, lightId, nodeDef });
-            lightTypesBound[nodeName] = lightId;
-            mx.HwShaderGenerator.bindLightShader(nodeDef, lightId++, genContext);
-        }
-
-        const lightDirection = light.getValueElement("direction").getValue().getData().data();
-        const lightColor = light.getValueElement("color").getValue().getData().data();
-        const lightIntensity = light.getValueElement("intensity").getValue().getData();
-
-        let rotatedLightDirection = new THREE.Vector3(...lightDirection)
-        rotatedLightDirection.transformDirection(getLightRotation())
-
-        lightData.push({
-            type: lightTypesBound[nodeName],
-            direction: rotatedLightDirection,
-            color: new THREE.Vector3(...lightColor),
-            intensity: lightIntensity,
-        });
+// Converts Three.js light type to MaterialX node name
+function threeLightTypeToMaterialXNodeName(threeLightType) {
+    switch (threeLightType) {
+        case 'PointLight':
+            return 'ND_point_light';
+        case 'DirectionalLight':
+            return 'ND_directional_light';
+        case 'SpotLight':
+            return 'ND_spot_light';
+        default:
+            console.warn('MaterialX: Unsupported light type: ' + threeLightType);
+            return 'ND_point_light'; // Default to point light
     }
+};
 
-    const threeLightTypeToMaterialXNodeName = (threeLightType) => {
-        switch (threeLightType) {
-            case 'PointLight':
-                return 'ND_point_light';
-            case 'DirectionalLight':
-                return 'ND_directional_light';
-            case 'SpotLight':
-                return 'ND_spot_light';
-            default:
-                console.warn('MaterialX: Unsupported light type: ' + threeLightType);
-                return 'ND_point_light'; // Default to point light
-        }
-    };
-
-    if (debug) console.log("Registering lights in MaterialX context", lights, lightData);
+/**
+ * Update light data for shader uniforms
+ * @param {Object} mx MaterialX Module
+ * @param {Array.<mx.Node>} lights Light nodes
+ * @param {mx.GenContext} genContext Shader generation context
+ * @returns {{ lightData: Array<any>, lightCount: number }}
+ */
+export function getLightData(mx, lights, genContext)
+{
+    const lightData = [];
+    const maxLightCount = genContext.getOptions().hwMaxActiveLightSources;
 
     // Three.js lights
     for (let light of lights) {
@@ -408,24 +396,51 @@ export async function registerLights(mx, lights, genContext)
         const lightDefinitionName = threeLightTypeToMaterialXNodeName(light.type);
 
         if (!lightTypesBound[lightDefinitionName])
-        {
-            lightTypesBound[lightDefinitionName] = lightId;
-            const nodeDef = null;
-            mx.HwShaderGenerator.bindLightShader(nodeDef, lightId++, genContext);
-        }
+            console.error("MaterialX: Light type not registered in context. Make sure to register light types before using them.", lightDefinitionName);
+
+        const wp = light.getWorldPosition(new THREE.Vector3());
+        const wd = getWorldDirection(light, new THREE.Vector3(0,0,-1));
+
+        // console.log("Registering light", light.penumbra);
 
         lightData.push({
             type: lightTypesBound[lightDefinitionName],
-            direction: light.direction?.clone() || new THREE.Vector3(0, -1, 0),
-            color: new THREE.Vector3().fromArray(light.color.toArray()),
-            intensity: light.intensity,
+            position: wp.clone(),
+            direction: wd.clone(),
+            color: new THREE.Color().fromArray(light.color.toArray()),
+            // Luminous efficacy for converting radiant power in watts (W) to luminous flux in lumens (lm) at a wavelength of 555 nm.
+            // Also, three.js lights don't have PI scale baked in, but MaterialX does, so we need to divide by PI for point and spot lights.
+            intensity: light.intensity * (light.isPointLight ? 683.0 / 3.1415 : light.isSpotLight ? 683.0 / 3.1415: 1.0),
+            decay_rate: 2.0,
+            // Approximations for testing – the relevant light has 61.57986...129.4445 as inner/outer spot angle
+            inner_angle: 0.9,
+            outer_angle: 0.4,
         });
     }
 
-    // Make sure max light count is large enough
-    genContext.getOptions().hwMaxActiveLightSources = Math.max(genContext.getOptions().hwMaxActiveLightSources, lightData.length);
+    // Count the number of lights that are not empty
+    const lightCount = lightData.length;
+
+    // If we don't have enough entries in lightData, fill with empty lights
+    while (lightData.length < maxLightCount) 
+    {
+        const emptyLight = {
+            type: 0, // Default light type
+            position: new THREE.Vector3(0, 0, 0),
+            direction: new THREE.Vector3(0, 0, -1),
+            color: new THREE.Color(0, 0, 0),
+            intensity: 0.0,
+            decay_rate: 2.0,
+            inner_angle: 0.0,
+            outer_angle: 0.0,
+        };
+        lightData.push(emptyLight);
+    }
+
+    if (debug) 
+        console.log("Registered lights in MaterialX context", lightTypesBound, lightData);
 
-    return lightData;
+    return { lightData, lightCount };
 }
 
 /**
@@ -435,11 +450,12 @@ export async function registerLights(mx, lights, genContext)
  */
 export function getUniformValues(shaderStage, textureLoader, searchPath, flipY)
 {
-    let threeUniforms = {};
+    const threeUniforms = {};
 
     const uniformBlocks = Object.values(shaderStage.getUniformBlocks());
     uniformBlocks.forEach(uniforms =>
     {
+        // TODO Seems struct uniforms (like in LightData) end up here as well, we should filter those out.
         if (!uniforms.empty())
         {
             for (let i = 0; i < uniforms.size(); ++i)
@@ -447,6 +463,7 @@ export function getUniformValues(shaderStage, textureLoader, searchPath, flipY)
                 const variable = uniforms.get(i);
                 const value = variable.getValue()?.getData();
                 const name = variable.getVariable();
+                if (debug) console.log("Adding uniform", { name, value, type: variable.getType().getName() });
                 threeUniforms[name] = new THREE.Uniform(toThreeUniform(variable.getType().getName(), value, name, uniforms,
                     textureLoader, searchPath, flipY));
             }

package/src/index.ts

@@ -1,5 +1,5 @@
-import { initializeMaterialX, state } from "./materialx.js";
+import { ready, state } from "./materialx.js";
 
 const getMaterialXEnvironment = () => state.materialXEnvironment;
 
-export { initializeMaterialX, getMaterialXEnvironment };
+export { ready, getMaterialXEnvironment };

package/src/loader/loader.needle.ts

@@ -38,23 +38,15 @@ export class MaterialXUniformUpdate extends Component {
 
         const camera = this.context.mainCamera;
         if (!camera) return;
-
+        
         MaterialXUniformUpdate.updateMaterial(material, gameObject, camera);
-
-        // If this is a Group, we need to update all direct children
-        if ((gameObject as any as Group).isGroup) {
-            gameObject.children.forEach((child: Object3D) => {
-                if (child instanceof Mesh && child.material)
-                    MaterialXUniformUpdate.updateMaterial(child.material, child, camera);
-            });
-        }
     }
 }
 
 export class MaterialXLoaderPlugin implements INeedleGLTFExtensionPlugin {
-    name = "MaterialXLoaderPlugin";
+    readonly name = "MaterialXLoaderPlugin";
 
-    mtlxLoader: MaterialXLoader | null = null;
+    private loader: MaterialXLoader | null = null;
 
     onImport = (loader: GLTFLoader, url: string, context: Context) => {
         if (debug) console.log("MaterialXLoaderPlugin: Registering MaterialX extension for", url);
@@ -62,13 +54,13 @@ export class MaterialXLoaderPlugin implements INeedleGLTFExtensionPlugin {
         // Register the MaterialX loader extension
         // Environment initialization is now handled in the MaterialXLoader constructor
         loader.register(p => {
-            this.mtlxLoader = new MaterialXLoader(p, context);
-            return this.mtlxLoader;
+            this.loader = new MaterialXLoader(p, context);
+            return this.loader;
         });
     };
 
-    onLoaded = (url: string, gltf: GLTF, _context: Context) => {
-        if (debug) console.log("MaterialXLoaderPlugin: glTF loaded", url, gltf.scene);
+    onLoaded = (url: string, gltf: GLTF, context: Context) => {
+        if (debug) console.log("[MaterialX] MaterialXLoaderPlugin: glTF loaded", url, gltf.scene);
 
         // Set up onBeforeRender callbacks for objects with MaterialX materials
         // This ensures uniforms are updated properly during rendering
@@ -78,25 +70,24 @@ export class MaterialXLoaderPlugin implements INeedleGLTFExtensionPlugin {
                 const material = mesh.material as Material;
 
                 if (material?.userData?.updateUniforms) {
-                    if (debug) console.log("Adding MaterialX uniform update component to:", child.name);
+                    if (debug) console.log("[MaterialX] Adding MaterialX uniform update component to:", child.name);
                     child.addComponent(MaterialXUniformUpdate);
                 }
             }
         });
 
-        if (debug) console.log("Loaded: ", this.mtlxLoader);
+        if (debug) console.log("[MaterialX] Loaded: ", this.loader);
 
         // Initialize MaterialX lighting system with scene data
         const environment = state.materialXEnvironment;
-        environment.initializeFromContext().then(() => {
-            if (this.mtlxLoader) {
-                this.mtlxLoader.updateLightingFromEnvironment(environment);
-            }
+        environment.initializeFromContext(context).then(() => {
+            console.warn("[MaterialX] Environment initialized...");
+            this.loader?.updateLightingFromEnvironment(environment);
         });
     };
 
     onExport = (_exporter: GLTFExporter, _context: Context) => {
-        console.log("TODO: MaterialXLoaderPlugin: Setting up export extensions");
+        console.log("[MaterialX] TODO: MaterialXLoaderPlugin: Setting up export extensions");
         // TODO: Add MaterialX export functionality if needed
     };
 }

package/src/loader/loader.three.ts

@@ -1,91 +1,59 @@
-import { Context } from "@needle-tools/engine";
-import { RawShaderMaterial, Material, MeshStandardMaterial, LoadingManager, TextureLoader, Texture, NearestFilter, Matrix4, GLSL3, AddEquation, OneMinusSrcAlphaFactor, SrcAlphaFactor, DoubleSide, Matrix3, Vector3, Object3D, Camera } from "three";
+import { Context, GameObject } from "@needle-tools/engine";
+import { ShaderMaterial, Material, MeshStandardMaterial, LoadingManager, TextureLoader, Texture, NearestFilter, Matrix4, GLSL3, AddEquation, OneMinusSrcAlphaFactor, SrcAlphaFactor, DoubleSide, Matrix3, Vector3, Object3D, Camera, Uniform } from "three";
 import { GLTFLoaderPlugin, GLTFParser } from "three/examples/jsm/loaders/GLTFLoader.js";
 import { getUniformValues } from "../helper.js";
-import { initializeMaterialX, MaterialXEnvironment, state } from "../materialx.js";
+import { ready, MaterialXEnvironment, state } from "../materialx.js";
 import { debug } from "../utils.js";
 
 // TypeScript interfaces matching the C# data structures
-interface MaterialXData {
-    version: string; // e.g. "1.39"
-    name: string;    // e.g. "Material"
-    mtlx: string;    // MaterialX XML content
+interface MaterialX_root_extension {
+    /** e.g. 1.39 */
+    version: string;
+    /** e.g. "Material" */
+    name: string;
+    /** MaterialX xml content */
+    mtlx: string;
 }
 
-interface MaterialXDataIndex {
-    name: string;    // Material name reference
+interface MaterialX_material_extension {
+    name: string; // Material name reference
 }
 
 // MaterialX loader extension for js GLTFLoader
 export class MaterialXLoader implements GLTFLoaderPlugin {
     name = "NEEDLE_materials_mtlx";
 
-    private rootMaterialXData: MaterialXData | null = null;
-    private parsedDocument: any = null;
-    private documentParsePromise: Promise<any> | null = null;
-    private rootDataInitialized = false;
+    // private rootMaterialXData: MaterialX_root_extension | null = null;
+    private _documentReadyPromise: Promise<any> | null = null;
     private environmentInitialized = false;
-    private generatedMaterials: RawShaderMaterial[] = [];
+    private generatedMaterials: ShaderMaterial[] = [];
 
-    constructor(private parser: GLTFParser, private context: Context) {
-        if (debug) console.log("MaterialXLoader created for parser");
-        // Initialize MaterialX environment after MaterialX is ready
-        this.initializeEnvironment();
-    }
-
-    // Initialize MaterialX environment - called once after MaterialX is ready
-    private async initializeEnvironment(): Promise<void> {
-        if (this.environmentInitialized) return;
-
-        if (debug) console.log("MaterialXLoader: Initializing MaterialX environment...");
-
-        // Ensure MaterialX is initialized first
-        await initializeMaterialX();
-
-        // Set up environment with context
-        const environment = state.materialXEnvironment;
-        environment.setContext(this.context);
-
-        // Initialize the environment from context (properly awaited)
-        try {
-            await environment.initializeFromContext();
-            this.environmentInitialized = true;
-            if (debug) console.log("MaterialXLoader: Environment initialized successfully");
-        } catch (error) {
-            console.warn("MaterialXLoader: Failed to initialize MaterialX environment:", error);
-        }
+    get materialX_root_data() {
+        return this.parser.json.extensions?.[this.name] as MaterialX_root_extension | null;
     }
 
-    // Initialize root data from parser.json.extensions once
-    private initializeRootData(): void {
-        if (this.rootDataInitialized) return;
-
-        const gltfExtensions = this.parser.json.extensions;
-        if (gltfExtensions?.[this.name]) {
-            if (debug) console.log("MaterialX extension found in root:", gltfExtensions[this.name]);
-
-            const materialXExtension = gltfExtensions[this.name];
-            this.rootMaterialXData = materialXExtension as MaterialXData;
+    constructor(private parser: GLTFParser, private context: Context) {
+        if (debug) console.log("MaterialXLoader created for parser");
+        // Start loading of MaterialX environment if the root extension exists
+        const hasMaterialXExtension = this.parser.json.extensions?.[this.name] != null;
+        if (hasMaterialXExtension) {
+            ready();
         }
-        this.rootDataInitialized = true;
     }
 
     // Parse the MaterialX document once and cache it
-    private async parseRootDocument(): Promise<any> {
-        if (this.documentParsePromise) {
-            return this.documentParsePromise;
+    private async _materialXDocumentReady(): Promise<any> {
+        if (this._documentReadyPromise) {
+            return this._documentReadyPromise;
         }
-
-        this.documentParsePromise = (async () => {
-            if (this.parsedDocument) return this.parsedDocument;
-
-            if (debug) console.log("Parsing MaterialX root document...");
+        return this._documentReadyPromise = (async () => {
+            if (debug) console.log("[MaterialX] Parsing MaterialX root document...");
 
             // Ensure MaterialX is initialized
-            await initializeMaterialX();
+            await ready();
 
             if (!state.materialXModule) {
-                throw new Error("MaterialX module failed to initialize");
+                throw new Error("[MaterialX] module failed to initialize");
             }
 
             // Create MaterialX document and parse ALL the XML data from root
@@ -93,18 +61,15 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
             doc.setDataLibrary(state.materialXStdLib);
 
             // Parse all MaterialX XML strings from the root data
-            if (this.rootMaterialXData) {
-                if (debug) console.log(`Parsing MaterialX XML for: ${this.rootMaterialXData.name}`);
-                await state.materialXModule.readFromXmlString(doc, this.rootMaterialXData.mtlx, "");
+            const root = this.materialX_root_data
+            if (root) {
+                if (debug) console.log(`[MaterialX] Parsing XML for: ${root.name}`);
+                await state.materialXModule.readFromXmlString(doc, root.mtlx, "");
             }
 
-            if (debug) console.log("MaterialX root document parsed successfully");
-
-            this.parsedDocument = doc;
+            if (debug) console.log("[MaterialX] root document parsed successfully");
             return doc;
         })();
-
-        return this.documentParsePromise;
     }
 
     loadMaterial(materialIndex: number): Promise<Material> | null {
@@ -118,14 +83,12 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
     }
 
     private async _loadMaterialAsync(materialIndex: number): Promise<Material> {
-        // Initialize root data first
-        this.initializeRootData();
 
         const materialDef = this.parser.json.materials?.[materialIndex];
-        if (debug) console.log("MaterialX extension found in material:", materialDef.extensions[this.name]);
+        if (debug) console.log("[MaterialX] extension found in material:", materialDef.extensions[this.name]);
 
         // Handle different types of MaterialX data
-        const dataIndex = materialDef.extensions[this.name] as MaterialXDataIndex;
+        const dataIndex = materialDef.extensions[this.name] as MaterialX_material_extension;
 
         if (dataIndex) {
             // Create a new material and process MaterialX - AWAIT THIS!
@@ -138,21 +101,14 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
         return fallbackMaterial;
     }
 
-    private rootDocument: Promise<any> | null = null;
-    private async createMaterialXMaterial(materialXData: MaterialXDataIndex): Promise<Material> {
+    private async createMaterialXMaterial(materialXData: MaterialX_material_extension): Promise<Material> {
         try {
             if (debug) console.log(`Creating MaterialX material: ${materialXData.name}`);
 
-            // Ensure MaterialX is initialized and document is parsed
-            await initializeMaterialX();
-
-            if (!this.rootDocument) {
-                this.rootDocument = this.parseRootDocument();
-            }
-            const doc = await this.rootDocument;
+            const doc = await this._materialXDocumentReady();
 
             if (!state.materialXModule || !state.materialXGenerator || !state.materialXGenContext) {
-                console.warn("MaterialX WASM module not ready, returning fallback material");
+                console.warn("[MaterialX] WASM module not ready, returning fallback material");
                 const fallbackMaterial = new MeshStandardMaterial();
                 fallbackMaterial.userData.materialX = materialXData;
                 fallbackMaterial.name = `MaterialX_Fallback_${materialXData.name}`;
@@ -163,11 +119,11 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
             let renderableElement = null;
             let foundRenderable = false;
 
-            if (debug) console.log("Mtlx doc", doc);
+            if (debug) console.log("[MaterialX] document", doc);
 
             // Search for material nodes first (following the reference pattern)
             const materialNodes = doc.getMaterialNodes();
-            if (debug) console.log(`Found ${materialNodes.length} material nodes in document`, materialNodes);
+            if (debug) console.log(`[MaterialX] Found ${materialNodes.length} material nodes in document`, materialNodes);
 
             // Handle both array and vector-like APIs
             const materialNodesLength = materialNodes.length;
@@ -175,13 +131,13 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
                 const materialNode = materialNodes[i];
                 if (materialNode) {
                     const materialName = materialNode.getNamePath();
-                    if (debug) console.log('Scan material: ', i, materialName);
+                    if (debug) console.log('[MaterialX] Scan material: ', i, materialName);
 
                     // Find the matching material
                     if (materialName == materialXData.name) {
                         renderableElement = materialNode;
                         foundRenderable = true;
-                        if (debug) console.log('-- add material: ', materialName);
+                        if (debug) console.log('[MaterialX] -- add material: ', materialName);
                         break;
                     }
                 }
@@ -235,35 +191,35 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
             */
 
             if (!renderableElement) {
-                console.warn("No renderable element found in MaterialX document");
+                console.warn("[MaterialX] No renderable element found in MaterialX document");
                 const fallbackMaterial = new MeshStandardMaterial();
                 fallbackMaterial.userData.materialX = materialXData;
                 fallbackMaterial.name = `MaterialX_NoRenderable_${materialXData.name}`;
                 return fallbackMaterial;
             }
 
-            if (debug) console.log("Using renderable element for shader generation");
+            if (debug) console.log("[MaterialX] Using renderable element for shader generation");
 
             // Check transparency and set context options like the reference
-            let isTransparent = state.materialXModule.isTransparentSurface(renderableElement, state.materialXGenerator.getTarget());
+            const isTransparent = state.materialXModule.isTransparentSurface(renderableElement, state.materialXGenerator.getTarget());
             state.materialXGenContext.getOptions().hwTransparency = isTransparent;
-            state.materialXGenContext.getOptions().shaderInterfaceType = state.materialXModule.ShaderInterfaceType.SHADER_INTERFACE_COMPLETE;
-            state.materialXGenContext.getOptions().hwSrgbEncodeOutput = true; // Like the reference
 
             // Generate shaders using the element's name path
-            if (debug) console.log("Generating MaterialX shaders...");
+            if (debug) console.log("[MaterialX] Generating MaterialX shaders...");
             const elementName = (renderableElement as any).getNamePath ? (renderableElement as any).getNamePath() : (renderableElement as any).getName();
 
             const shader = state.materialXGenerator.generate(elementName, renderableElement, state.materialXGenContext);
 
-            // Get vertex and fragment shader source
-            // Remove #version directive for newer js. It's added by RawShaderMaterial glslVersion.
+            // Get vertex and fragment shader source, and remove #version directive for newer js. 
+            // It's added by three.js glslVersion.
             let vertexShader = shader.getSourceCode("vertex").replace(/^#version.*$/gm, '').trim();
             let fragmentShader = shader.getSourceCode("pixel").replace(/^#version.*$/gm, '').trim();
 
             // MaterialX uses different attribute names than js defaults,
             // so we patch the MaterialX shaders to match the js standard names.
-            // Otherwise, we'd have to modify the mesh attributes (see below).
+            // Otherwise, we'd have to modify the mesh attributes (see original MaterialX for reference).
+
+            // Patch vertexShader
             vertexShader = vertexShader.replace(/\bi_position\b/g, 'position');
             vertexShader = vertexShader.replace(/\bi_normal\b/g, 'normal');
             vertexShader = vertexShader.replace(/\bi_texcoord_0\b/g, 'uv');
@@ -271,6 +227,7 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
             vertexShader = vertexShader.replace(/\bi_tangent\b/g, 'tangent');
             vertexShader = vertexShader.replace(/\bi_color_0\b/g, 'color');
 
+            // Patch fragmentShader
             fragmentShader = fragmentShader.replace(/\bi_position\b/g, 'position');
             fragmentShader = fragmentShader.replace(/\bi_normal\b/g, 'normal');
             fragmentShader = fragmentShader.replace(/\bi_texcoord_0\b/g, 'uv');
@@ -278,19 +235,39 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
             fragmentShader = fragmentShader.replace(/\bi_tangent\b/g, 'tangent');
             fragmentShader = fragmentShader.replace(/\bi_color_0\b/g, 'color');
 
-            // From the original code:
-
-            /* PATCHING THE SHADER IS BETTER – THIS IS JUST FOR REFERENCE 
-            // Use default MaterialX naming convention.
-            var startStreamTime = performance.now();
-            child.geometry.attributes.i_position = child.geometry.attributes.position;
-            child.geometry.attributes.i_normal = child.geometry.attributes.normal;
-            child.geometry.attributes.i_tangent = child.geometry.attributes.tangent;
-            child.geometry.attributes.i_texcoord_0 = child.geometry.attributes.uv;
-            */
+            // Remove `in vec3 position;` and so on since they're already declared by ShaderMaterial
+            vertexShader = vertexShader.replace(/in\s+vec3\s+position;/g, '');
+            vertexShader = vertexShader.replace(/in\s+vec3\s+normal;/g, '');
+            vertexShader = vertexShader.replace(/in\s+vec3\s+uv;/g, '');
+            vertexShader = vertexShader.replace(/in\s+vec3\s+uv1;/g, '');
+            vertexShader = vertexShader.replace(/in\s+vec4\s+tangent;/g, '');
+            vertexShader = vertexShader.replace(/in\s+vec4\s+color;/g, '');
+
+            // Patch uv 2-component to 3-component (`texcoord_0 = uv;` needs to be replaced with `texcoord_0 = vec3(uv, 0.0);`)
+            // TODO what if we actually have a 3-component UV? Not sure what three.js does then
+            vertexShader = vertexShader.replace(/texcoord_0 = uv;/g, 'texcoord_0 = vec3(uv, 0.0);');
+
+            // Patch units – seems MaterialX uses different units and we end up with wrong light values?
+            // result.direction = light.position - position;
+            fragmentShader = fragmentShader.replace(
+                /result\.direction\s*=\s*light\.position\s*-\s*position;/g,
+                'result.direction = (light.position - position) * 10.0 / 1.0;');
+
+            // Add tonemapping and colorspace handling
+            // Replace `out vec4 out1;` with `out vec4 gl_FragColor;`
+            fragmentShader = fragmentShader.replace(
+                /out\s+vec4\s+out1;/,
+                'layout(location = 0) out vec4 pc_fragColor;\n#define gl_FragColor pc_fragColor');
+
+            // Replace `out1 = vec4(<CAPTURE>)` with `gl_FragColor = vec4(<CAPTURE>)` and tonemapping/colorspace handling
+            fragmentShader = fragmentShader.replace(/^\s*out1\s*=\s*vec4\((.*)\);/gm,
+                `
+    gl_FragColor = vec4($1);
+    #include <tonemapping_fragment>
+    #include <colorspace_fragment>`);
 
             if (debug) {
-                console.group("Material: ", materialXData.name);
+                console.group("[MaterialX]: ", materialXData.name);
                 console.log("Vertex shader length:", vertexShader.length, vertexShader);
                 console.log("Fragment shader length:", fragmentShader.length, fragmentShader);
                 console.groupEnd();
@@ -303,7 +280,8 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
                 // Find the texture in the textures in the parser and load it from there
                 return url;
             });
-            let textureLoader = new TextureLoader(loadingManager);
+
+            const textureLoader = new TextureLoader(loadingManager);
 
             // Override the textureLoader.load method to use the parser's loadTexture directly,
             // since we want to load the textures from the glTF document and not from disk.
@@ -329,17 +307,17 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
                     const textures = ext?.textures || [];
 
                     const index = textures.findIndex(tex => {
-                        if (debug) console.log("Checking texture:", tex.name, "against URL:", filenameWithoutExt);
+                        if (debug) console.log("[MaterialX] Checking texture:", tex.name, "against URL:", filenameWithoutExt);
                         return tex.name === filenameWithoutExt;
                     });
 
                     if (index < 0) {
-                        console.warn("Texture not found in parser:", filenameWithoutExt, this.parser.json);
+                        console.warn("[MaterialX] Texture not found in parser:", filenameWithoutExt, this.parser.json);
                         onError?.(new Error(`Texture not found: ${filenameWithoutExt}`));
                         return;
                     }
                     this.parser.getDependency("texture", index).then(tex => {
-                        if (debug) console.log("Texture loaded:", tex);
+                        if (debug) console.log("[MaterialX] Texture loaded:", tex);
                         // update the checkerboard texture with the loaded texture
                         checkerboardTexture.image = tex.image;
                         checkerboardTexture.needsUpdate = true;
@@ -355,7 +333,7 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
             const searchPath = ""; // Could be derived from the asset path if needed
             const flipV = false; // Set based on your geometry requirements
 
-            let uniforms = {
+            const uniforms = {
                 ...getUniformValues(shader.getStage('vertex'), textureLoader, searchPath, flipV),
                 ...getUniformValues(shader.getStage('pixel'), textureLoader, searchPath, flipV),
             };
@@ -371,20 +349,12 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
                 return new Matrix4();
             };
 
-            if (debug) console.log({ lightData, radianceTexture, irradianceTexture });
-
-            Object.assign(uniforms, {
-                u_numActiveLightSources: { value: lightData?.length || 0, type: 'i' },
-            });
+            if (debug) console.log("Lights", { lightData, radianceTexture, irradianceTexture });
 
-            if (lightData?.length) {
-                Object.assign(uniforms, {
-                    u_lightData: { value: lightData, type: 'v4v' },
-                });
+            if (debug) {
+                const mips = Math.trunc(Math.log2(Math.max(radianceTexture?.width ?? 0, radianceTexture?.height ?? 0))) + 1;
+                console.log("[MaterialX] Radiance texture mips:", mips, "for texture size:", radianceTexture?.width, "x", radianceTexture?.height);
             }
-
-            const mips = Math.trunc(Math.log2(Math.max(radianceTexture?.width ?? 0, radianceTexture?.height ?? 0))) + 1;
-            if (debug) console.log("Radiance texture mips:", mips, "for texture size:", radianceTexture?.width, "x", radianceTexture?.height);
             Object.assign(uniforms, {
                 u_envMatrix: { value: getLightRotation() },
                 u_envRadiance: { value: radianceTexture, type: 't' },
@@ -393,6 +363,7 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
                 u_envRadianceSamples: { value: 8, type: 'i' },
                 u_envIrradiance: { value: irradianceTexture, type: 't' },
                 u_refractionEnv: { value: true },
+                u_lightData: { value: [] },
             });
 
             // console.log("NEW MATERIAL UNIFORMS", uniforms);
@@ -402,7 +373,7 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
             // console.log("Generated fragment shader:", fragmentShader.substring(0, 500) + "...");
 
             // Create js RawShaderMaterial (with GLSL3 for MaterialX compatibility)
-            const shaderMaterial = new RawShaderMaterial({
+            const shaderMaterial = new ShaderMaterial({
                 uniforms: uniforms,
                 vertexShader: vertexShader,
                 fragmentShader: fragmentShader,
@@ -466,14 +437,32 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
                     uniforms.u_frame.value = time;
                 }
 
+                // Update light uniforms
+                const environment = state.materialXEnvironment;
+                const lightData = environment.getLightData() || null;
+                const lightCount = environment.getLightCount() || 0;
+                if (uniforms.u_numActiveLightSources) {
+                    uniforms.u_numActiveLightSources.value = lightCount;
+                }
+
+                if (lightData && uniforms.u_lightData) {
+                    uniforms.u_lightData.value = lightData;
+                    if (debug) console.log("Updating light data for material", shaderMaterial.name, lightData, shaderMaterial.uniforms);
+                }
+
                 // Mark uniforms as needing update
                 shaderMaterial.uniformsNeedUpdate = true;
             };
 
+            this.context.pre_update_callbacks.push(() => {
+                const environment = state.materialXEnvironment;
+                environment.updateLighting(false);
+            });
+
             shaderMaterial.name = `MaterialX_Generated_${materialXData.name}`;
 
             // Add debugging to see if the material compiles correctly
-            if (debug) console.log("MaterialX material created successfully:", shaderMaterial.name);
+            if (debug) console.log("[MaterialX]  material created successfully:", shaderMaterial.name);
             // if (debug) console.log("Material uniforms keys:", Object.keys(shaderMaterial.uniforms || {}));
             // if (debug) console.log("Material transparent:", shaderMaterial.transparent);
             // if (debug) console.log("Material side:", shaderMaterial.side);
@@ -485,7 +474,7 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
 
         } catch (error) {
             // This is a wasm error (an int) that we need to resolve
-            console.error("Error creating MaterialX material:", error);
+            console.error("[MaterialX] Error creating MaterialX material:", error);
             // Return a fallback material with stored MaterialX data
             const fallbackMaterial = new MeshStandardMaterial();
             fallbackMaterial.userData.materialX = materialXData;
@@ -499,31 +488,35 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
     updateLightingFromEnvironment(environment: MaterialXEnvironment): void {
 
         // Get lighting data from environment
-        // const lights = environment.getLights() || [];
         const lightData = environment.getLightData() || null;
+        const lightCount = environment.getLightCount() || 0;
         const radianceTexture = environment.getRadianceTexture() || null;
         const irradianceTexture = environment.getIrradianceTexture() || null;
 
-        if (debug) console.log(`Updating lighting for ${this.generatedMaterials.length} MaterialX materials`, {
-            lightData, radianceTexture, irradianceTexture,
-        });
+        if (debug) {
+            console.log(`[MaterialX] Updating lighting for ${this.generatedMaterials.length} MaterialX materials`, {
+                lightData, radianceTexture, irradianceTexture,
+            });
+        }
 
         // Update each generated material's lighting uniforms
         this.generatedMaterials.forEach((material, _index) => {
             if (!material.uniforms) return;
 
+            console.warn(material.name, material.uniforms, lightCount)
+
             // Update light count
-            if (material.uniforms.u_numActiveLightSources && lightData) {
-                material.uniforms.u_numActiveLightSources.value = lightData.length;
+            if (material.uniforms.u_numActiveLightSources && lightCount >= 0) {
+                material.uniforms.u_numActiveLightSources.value = lightCount;
             }
 
             // Update light data if we have lights
             if (lightData) {
-                if (!material.uniforms.u_lightData) {
-                    material.uniforms.u_lightData = { value: null };
-                }
+                material.uniforms.u_lightData ??= new Uniform(null)
                 material.uniforms.u_lightData.value = lightData;
+                if (debug) console.log("[MaterialX] Updated light data for material", material.name, lightData, material.uniforms,);
             }
+            else if (debug) console.warn("[MaterialX] No light data available to update uniforms for material", material.name);
 
             // Update environment uniforms
             if (material.uniforms.u_envMatrix) {
@@ -542,14 +535,7 @@ export class MaterialXLoader implements GLTFLoaderPlugin {
             // Mark uniforms as needing update
             // console.log("Light data in uniforms", material.uniforms, material.fragmentShader);
             material.uniformsNeedUpdate = true;
+            console.log(material)
         });
-
-        /* TODO not working yet
-        this.generatedMaterials.forEach((material, index) => {
-            console.log("Regenerating shaders for MaterialX material:", index, material.name);
-            material.userData.regenerateShaders();
-            resetShaders(this.context);
-        });
-        */
     }
 }

package/src/materialx.ts

@@ -1,14 +1,13 @@
-import { Context, delay, isDevEnvironment, ObjectUtils } from "@needle-tools/engine";
+import { Context, delay, isDevEnvironment, ObjectUtils, GameObject } from "@needle-tools/engine";
 import MaterialX from "../bin/JsMaterialXGenShader.js";
 import { debug } from "./utils.js";
 import { renderPMREMToEquirect } from "./textureHelper.js";
-import { Light, MeshBasicMaterial, Object3D, PMREMGenerator } from "three";
-import { registerLights } from "./helper.js";
+import { Light, MeshBasicMaterial, Object3D, PMREMGenerator, Texture } from "three";
+import { registerLights, getLightData } from "./helper.js";
 
 
-// Global MaterialX module instance - initialized lazily
 export const state = new class {
-    materialXModule: any = null;
+    materialXModule: typeof MaterialX | null = null;
     materialXGenerator: any = null;
     materialXGenContext: any = null;
     materialXStdLib: any = null;
@@ -22,14 +21,14 @@ export const state = new class {
     }
 }
 
-// Initialize MaterialX WASM module lazily
-export async function initializeMaterialX(): Promise<void> {
+/** Initialize the MaterialX module. Must be awaited before trying to create materials */
+export async function ready(): Promise<void> {
     if (state.materialXInitPromise) {
         return state.materialXInitPromise;
     }
     return state.materialXInitPromise = (async () => {
         if (state.materialXModule) return; // Already initialized
-        if (debug) console.log("Initializing MaterialX WASM module...");
+        if (debug) console.log("[MaterialX] Initializing WASM module...");
         try {
 
             const urls: Array<string> = await Promise.all([
@@ -44,7 +43,7 @@ export async function initializeMaterialX(): Promise<void> {
 
             const module = await MaterialX({
                 locateFile: (path: string, scriptDirectory: string) => {
-                    if (debug) console.debug("MaterialX locateFile called:", { path, scriptDirectory });
+                    if (debug) console.debug("[MaterialX]  locateFile called:", { path, scriptDirectory });
 
                     if (path.includes("JsMaterialXCore.wasm")) {
                         return JsMaterialXCore; // Use the URL for the core WASM file
@@ -59,7 +58,7 @@ export async function initializeMaterialX(): Promise<void> {
                     return scriptDirectory + path;
                 },
             });
-            if (debug) console.log("MaterialXLoader module loaded", module);
+            if (debug) console.log("[MaterialX] module loaded", module);
             state.materialXModule = module;
 
             // Initialize shader generator and context
@@ -71,23 +70,36 @@ export async function initializeMaterialX(): Promise<void> {
             state.materialXStdLib = module.loadStandardLibraries(state.materialXGenContext);
             tempDoc.setDataLibrary(state.materialXStdLib);
 
-            // Initialize basic lighting with default light rig
-            //             const defaultLightRigXml = `<?xml version="1.0"?>
-            // <materialx version="1.39">
-            //   <!-- Default directional light -->
-            //   <directional_light name="default_light" type="lightshader">
-            //     <input name="direction" type="vector3" value="0.0, -1.0, -0.5" />
-            //     <input name="color" type="color3" value="1.0, 1.0, 1.0" />
-            //     <input name="intensity" type="float" value="1.0" />
-            //   </directional_light>
-            // </materialx>`;
+            // TODO ShaderInterfaceType.SHADER_INTERFACE_REDUCED would be better, but doesn't actually seem to be supported in the MaterialX javascript bindings
+            state.materialXGenContext.getOptions().shaderInterfaceType = state.materialXModule.ShaderInterfaceType.SHADER_INTERFACE_COMPLETE;
+
+            // SPECULAR_ENVIRONMENT_NONE: Do not use specular environment maps.
+            // SPECULAR_ENVIRONMENT_FIS: Use Filtered Importance Sampling for specular environment/indirect lighting.
+            // SPECULAR_ENVIRONMENT_PREFILTER: Use pre-filtered environment maps for specular environment/indirect lighting.
+            state.materialXGenContext.getOptions().hwSpecularEnvironmentMethod = state.materialXModule.HwSpecularEnvironmentMethod.SPECULAR_ENVIRONMENT_FIS;
+
+            // TRANSMISSION_REFRACTION: Use a refraction approximation for transmission rendering.
+            // TRANSMISSION_OPACITY: Use opacity for transmission rendering.
+            // state.materialXGenContext.getOptions().hwTransmissionRenderMethod = state.materialXModule.HwTransmissionRenderMethod.TRANSMISSION_REFRACTION;
+
+            // Turned off because we're doing color space conversion the three.js way
+            state.materialXGenContext.getOptions().hwSrgbEncodeOutput = false;
+
+            // Enables the generation of a prefiltered environment map.
+            // TODO Would be great to use but requires setting more uniforms (like u_envPrefilterMip).
+            // When set to true, the u_envRadiance map is expected to be a prefiltered environment map.
+            // state.materialXGenContext.getOptions().hwWriteEnvPrefilter = true;
+
+            // Set a reasonable default for max active lights
+            state.materialXGenContext.getOptions().hwMaxActiveLightSources = 4;
 
             // This prewarms the shader generation context to have all light types
-            await registerLights(state.materialXModule, [], state.materialXGenContext);
+            await registerLights(state.materialXModule, state.materialXGenContext);
+            // getLightData(state.materialXModule, [], state.materialXGenContext);
 
-            if (debug) console.log("MaterialX generator initialized successfully");
+            if (debug) console.log("[MaterialX]  generator initialized successfully");
         } catch (error) {
-            console.error("Failed to load MaterialX module:", error);
+            console.error("[MaterialX] Failed to load MaterialX module:", error);
             throw error;
         }
     })();
@@ -95,134 +107,127 @@ export async function initializeMaterialX(): Promise<void> {
 
 // MaterialX Environment Manager - handles lighting and environment setup
 export class MaterialXEnvironment {
-    private lights: any[] = [];
-    private lightData: any = null;
-    private radianceTexture: any = null;
-    private irradianceTexture: any = null;
-    private context: Context | null = null;
-    private initialized: boolean = false;
+    private _context: Context | null = null;
+    private _lights: Array<Light> = [];
+    private _lightData: any = null;
+    private _lightCount: number = 0;
+    private _radianceTexture: Texture | null = null;
+    private _irradianceTexture: Texture | null = null;
+    private _initializePromise: Promise<boolean> | null = null;
 
     constructor() {
-        if (debug) console.log("MaterialX Environment created");
+        if (debug) console.log("[MaterialX] Environment created");
     }
 
-    setContext(context: Context) {
-        this.context = context;
+    // Initialize with Needle Engine context
+    async initializeFromContext(context: Context): Promise<boolean> {
+        if (this._initializePromise) {
+            return this._initializePromise;
+        }
+        return this._initializePromise = this._initialize(context);
     }
 
-    /*
-    // Initialize MaterialX lighting system based on the reference implementation
-    async initializeLighting(lightRigXml: string, renderer?: any, radianceTexture?: any, irradianceTexture?: any): Promise<void> {
-        if (!materialXModule || !materialXGenContext) {
-            console.warn("MaterialX module not initialized, skipping lighting setup");
-            return;
-        }
+    getLightData() { return this._lightData; }
+    getLightCount() { return this._lightCount; }
 
-        registerLights(materialXModule, this.lights, materialXGenContext);
-    }
-    */
+    setRadianceTexture(texture: Texture) { this._radianceTexture = texture; }
+    getRadianceTexture() { return this._radianceTexture; }
 
-    // Initialize with Needle Engine context
-    async initializeFromContext(): Promise<void> {
-        if (!this.context) {
-            console.warn("No Needle context available for MaterialX environment initialization");
-            return;
-        }
+    getIrradianceTexture() { return this._irradianceTexture; }
+    setIrradianceTexture(texture: Texture) { this._irradianceTexture = texture; }
 
-        // Prevent multiple initializations
-        if (this.initialized) {
-            if (debug) console.log("MaterialX environment already initialized, skipping");
-            return;
-        }
+    private _initialize: (context: Context) => Promise<boolean> = async (context: Context) => {
+
+        this._context = context;
+
+        // Get renderer from context
+        const renderer = context.renderer;
 
         // Clean up previous textures if they exist
-        if (this.radianceTexture) {
-            if (debug) console.log("Disposing previous radiance texture");
-            this.radianceTexture.dispose();
-            this.radianceTexture = null;
+        if (this._radianceTexture) {
+            if (debug) console.log("[MaterialX] Disposing previous radiance texture");
+            this._radianceTexture.dispose();
+            this._radianceTexture = null;
         }
-        if (this.irradianceTexture) {
-            if (debug) console.log("Disposing previous irradiance texture");
-            this.irradianceTexture.dispose();
-            this.irradianceTexture = null;
+        if (this._irradianceTexture) {
+            if (debug) console.log("[MaterialX] Disposing previous irradiance texture");
+            this._irradianceTexture.dispose();
+            this._irradianceTexture = null;
         }
 
-        // Get renderer from context
-        const renderer = this.context.renderer;
-
         // TODO remove this delay; we should wait for the scene lighting to be ready
         // and then update the uniforms
-        let envMap = this.context.scene.environment;
+        let envMap = context.scene.environment;
         while (!envMap) {
-            await delay(200);
-            envMap = this.context.scene.environment;
+            await delay(1000);
+            envMap = context.scene.environment;
         }
-        var pmrem = new PMREMGenerator(renderer);
+        const pmrem = new PMREMGenerator(renderer);
         const target = pmrem.fromEquirectangular(envMap);
-
         const radianceRenderTarget = renderPMREMToEquirect(renderer, target.texture, 0.0, 1024, 512, target.height);
         const irradianceRenderTarget = renderPMREMToEquirect(renderer, target.texture, 1.0, 32, 16, target.height);
-
-        this.radianceTexture = radianceRenderTarget.texture;
-        this.irradianceTexture = irradianceRenderTarget.texture;
+        this._radianceTexture = radianceRenderTarget.texture;
+        this._irradianceTexture = irradianceRenderTarget.texture;
 
         // Clean up PMREM generator and its render target
         target.dispose();
         pmrem.dispose();
 
         if (debug) {
-            console.log({ radiance: this.radianceTexture, irradiance: this.irradianceTexture });
+            console.log({ radiance: this._radianceTexture, irradiance: this._irradianceTexture });
             // Show both of them on cubes in the scene
             const unlitMat = new MeshBasicMaterial();
+            unlitMat.side = 2;
             const radianceMat = unlitMat.clone();
-            radianceMat.map = this.radianceTexture;
-            const radianceCube = ObjectUtils.createPrimitive("Cube", { material: radianceMat });
+            radianceMat.map = this._radianceTexture;
+            const radianceCube = ObjectUtils.createPrimitive("Quad", { material: radianceMat });
             const irradianceMat = unlitMat.clone();
-            irradianceMat.map = this.irradianceTexture;
-            const irradianceCube = ObjectUtils.createPrimitive("Cube", { material: irradianceMat });
-            this.context.scene.add(radianceCube);
-            this.context.scene.add(irradianceCube);
+            irradianceMat.map = this._irradianceTexture;
+            const irradianceCube = ObjectUtils.createPrimitive("Quad", { material: irradianceMat });
+            context.scene.add(radianceCube);
+            context.scene.add(irradianceCube);
             radianceCube.position.set(2, 0, 0);
-            radianceCube.scale.y = 0.00001;
             irradianceCube.position.set(-2, 0, 0);
-            irradianceCube.scale.y = 0.00001;
-            // await this.initializeLighting(defaultLightRigXml, renderer);
-            console.log("MaterialX environment initialized from Needle context", this, this.context.scene);
+            console.log("[MaterialX] environment initialized from Needle context", this, this._context.scene);
         }
 
-        // Find lights in scene
-        let lights = new Array<Light>();
-        this.context.scene.traverse((object: Object3D) => {
-            if ((object as Light).isLight) lights.push(object as Light);
-        });
-
-        this.lightData = await registerLights(state.materialXModule, lights, state.materialXGenContext);
+        this.updateLighting(true);
 
         // Mark as initialized
-        this.initialized = true;
+        return true;
     }
 
-    // getLights() { return this.lights; }
-    getLightData() { return this.lightData; }
-    getRadianceTexture() { return this.radianceTexture; }
-    getIrradianceTexture() { return this.irradianceTexture; }
+    updateLighting(collectLights: boolean) {
+        if (!this._context) return;
+
+        // Find lights in scene
+        if (collectLights) {
+            const lights = new Array<Light>();
+            this._context.scene.traverse((object: Object3D) => {
+                if ((object as Light).isLight && GameObject.isActiveInHierarchy(object))
+                    lights.push(object as Light);
+            });
+            this._lights = lights;
+        }
 
-    setRadianceTexture(texture: any) { this.radianceTexture = texture; }
-    setIrradianceTexture(texture: any) { this.irradianceTexture = texture; }
+        const { lightData, lightCount } = getLightData(state.materialXModule, this._lights, state.materialXGenContext);
+        this._lightData = lightData;
+        this._lightCount = lightCount;
+    }
 
     // Reset the environment to allow re-initialization
     reset() {
-        if (debug) console.log("Resetting MaterialX environment");
-        if (this.radianceTexture) {
-            this.radianceTexture.dispose();
-            this.radianceTexture = null;
+        if (debug) console.log("[MaterialX] Resetting environment");
+        if (this._radianceTexture) {
+            this._radianceTexture.dispose();
+            this._radianceTexture = null;
         }
-        if (this.irradianceTexture) {
-            this.irradianceTexture.dispose();
-            this.irradianceTexture = null;
+        if (this._irradianceTexture) {
+            this._irradianceTexture.dispose();
+            this._irradianceTexture = null;
         }
-        this.initialized = false;
-        // this.lights = [];
-        this.lightData = null;
+        this._initializePromise = null;
+        this._lights = [];
+        this._lightData = null;
     }
 }