render_path_tracing.cu

// ======================================================================== //
// Copyright (c) Meta Platforms, Inc. and affiliates.                       //
//                                                                          //
// This source code is licensed under the MIT license found in the          //
// LICENSE file in the root directory of this source tree.                  //
// ======================================================================== //

#include "render_path_tracing.h"
#include "render_task.h"
#include "utils.cuh"
#include <sys/stat.h>
#include <algorithm>
#include <iomanip>
#include <sstream>

// ====================================================
// SHARC Management Functions
// ====================================================
void RenderWindowPT::initializeSharc()
{
    if (sharcInitialized)
        destroySharc();
    
    cudaError_t err;
    
    err = cudaMalloc(&d_sharcHashEntries, sharcCapacity * sizeof(uint64_t));
    if (err != cudaSuccess) {
        printf("SHARC: Failed to allocate hash entries\n");
        return;
    }
    
    err = cudaMalloc(&d_sharcLockBuffer, sharcCapacity * sizeof(uint32_t));
    if (err != cudaSuccess) {
        cudaFree(d_sharcHashEntries); d_sharcHashEntries = nullptr;
        printf("SHARC: Failed to allocate lock buffer\n");
        return;
    }
    
    err = cudaMalloc(&d_sharcAccumBuffer, sharcCapacity * sizeof(SharcAccumulationData));
    if (err != cudaSuccess) {
        cudaFree(d_sharcHashEntries); d_sharcHashEntries = nullptr;
        cudaFree(d_sharcLockBuffer); d_sharcLockBuffer = nullptr;
        printf("SHARC: Failed to allocate accumulation buffer\n");
        return;
    }
    
    err = cudaMalloc(&d_sharcResolvedBuffer, sharcCapacity * sizeof(SharcPackedData));
    if (err != cudaSuccess) {
        cudaFree(d_sharcHashEntries); d_sharcHashEntries = nullptr;
        cudaFree(d_sharcLockBuffer); d_sharcLockBuffer = nullptr;
        cudaFree(d_sharcAccumBuffer); d_sharcAccumBuffer = nullptr;
        printf("SHARC: Failed to allocate resolved buffer\n");
        return;
    }

    err = cudaMalloc(&d_sharcHairMetaBuffer, sharcCapacity * sizeof(SharcHairMetaData));
    if (err != cudaSuccess) {
        cudaFree(d_sharcHashEntries); d_sharcHashEntries = nullptr;
        cudaFree(d_sharcLockBuffer); d_sharcLockBuffer = nullptr;
        cudaFree(d_sharcAccumBuffer); d_sharcAccumBuffer = nullptr;
        cudaFree(d_sharcResolvedBuffer); d_sharcResolvedBuffer = nullptr;
        printf("SHARC: Failed to allocate hair meta buffer\n");
        return;
    }
    
    cudaMemset(d_sharcHashEntries, 0, sharcCapacity * sizeof(uint64_t));
    cudaMemset(d_sharcLockBuffer, 0, sharcCapacity * sizeof(uint32_t));
    cudaMemset(d_sharcAccumBuffer, 0, sharcCapacity * sizeof(SharcAccumulationData));
    cudaMemset(d_sharcResolvedBuffer, 0, sharcCapacity * sizeof(SharcPackedData));
    cudaMemset(d_sharcHairMetaBuffer, 0, sharcCapacity * sizeof(SharcHairMetaData));
    
    sharcPrevCameraPos = make_float3(camera.getFrom().x, camera.getFrom().y, camera.getFrom().z);
    sharcInitialized = true;
    
    // Update launch params with buffer pointers AND capacity
    owlParamsSet1ul(this->launchParams, "sharcParams.hashMapData.hashEntriesBuffer", (uint64_t)d_sharcHashEntries);
    owlParamsSet1ul(this->launchParams, "sharcParams.hashMapData.lockBuffer", (uint64_t)d_sharcLockBuffer);
    owlParamsSet1ui(this->launchParams, "sharcParams.hashMapData.capacity", sharcCapacity);
    owlParamsSet1ul(this->launchParams, "sharcParams.accumulationBuffer", (uint64_t)d_sharcAccumBuffer);
    owlParamsSet1ul(this->launchParams, "sharcParams.resolvedBuffer", (uint64_t)d_sharcResolvedBuffer);
    owlParamsSet1ul(this->launchParams, "sharcParams.hairMetaBuffer", (uint64_t)d_sharcHairMetaBuffer);
    owlParamsSet1f(this->launchParams, "sharcParams.radianceScale", sharcRadianceScale);
    owlParamsSet1i(this->launchParams, "sharcParams.hairQueryMode", sharcHairQueryMode);
    owlParamsSet1f(this->launchParams, "sharcParams.hairTangentThreshold", sharcHairTangentThreshold);
    owlParamsSet1f(this->launchParams, "sharcParams.hairParallelScale", sharcHairParallelScale);
    owlParamsSet1f(this->launchParams, "sharcParams.hairPerpScale", sharcHairPerpScale);
    float effectiveSceneScale = (sharcSceneScale > 0.0f) ? sharcSceneScale : this->sceneScale;
    owlParamsSet1f(this->launchParams, "sharcParams.gridParameters.sceneScale", effectiveSceneScale);
    owlParamsSet1f(this->launchParams, "sharcParams.gridParameters.logarithmBase", this->sharcLogBase);
    owlParamsSet1f(this->launchParams, "sharcParams.gridParameters.levelBias", (float)this->sharcLevelBias);
    
    printf("SHARC: Initialized with %u entries (%.2f MB)\n", 
           sharcCapacity, 
           float(sharcCapacity * (sizeof(uint64_t) + sizeof(uint32_t) + 
                                  sizeof(SharcAccumulationData) + sizeof(SharcPackedData) +
                                  sizeof(SharcHairMetaData))) / (1024.0f * 1024.0f));
}

void RenderWindowPT::destroySharc()
{
    if (d_sharcHashEntries) { cudaFree(d_sharcHashEntries); d_sharcHashEntries = nullptr; }
    if (d_sharcLockBuffer) { cudaFree(d_sharcLockBuffer); d_sharcLockBuffer = nullptr; }
    if (d_sharcAccumBuffer) { cudaFree(d_sharcAccumBuffer); d_sharcAccumBuffer = nullptr; }
    if (d_sharcResolvedBuffer) { cudaFree(d_sharcResolvedBuffer); d_sharcResolvedBuffer = nullptr; }
    if (d_sharcHairMetaBuffer) { cudaFree(d_sharcHairMetaBuffer); d_sharcHairMetaBuffer = nullptr; }
    
    // Clear pointers in launch params
    owlParamsSet1ul(this->launchParams, "sharcParams.hashMapData.hashEntriesBuffer", 0);
    owlParamsSet1ul(this->launchParams, "sharcParams.hashMapData.lockBuffer", 0);
    owlParamsSet1ul(this->launchParams, "sharcParams.accumulationBuffer", 0);
    owlParamsSet1ul(this->launchParams, "sharcParams.resolvedBuffer", 0);
    owlParamsSet1ul(this->launchParams, "sharcParams.hairMetaBuffer", 0);
    
    sharcInitialized = false;
}

void RenderWindowPT::resetSharc()
{
    if (!sharcInitialized) return;
    
    cudaMemset(d_sharcHashEntries, 0, sharcCapacity * sizeof(uint64_t));
    cudaMemset(d_sharcLockBuffer, 0, sharcCapacity * sizeof(uint32_t));
    cudaMemset(d_sharcAccumBuffer, 0, sharcCapacity * sizeof(SharcAccumulationData));
    cudaMemset(d_sharcResolvedBuffer, 0, sharcCapacity * sizeof(SharcPackedData));
    cudaMemset(d_sharcHairMetaBuffer, 0, sharcCapacity * sizeof(SharcHairMetaData));
    printf("SHARC: Cache reset\n");
}

void RenderWindowPT::resolveSharc()
{
    if (!sharcInitialized || !d_sharcHashEntries || !d_sharcAccumBuffer || !d_sharcResolvedBuffer)
        return;
    
    // Clear every frame mode: reset resolved buffer to discard temporal history
    if (sharcClearEveryFrame) {
        cudaMemset(d_sharcResolvedBuffer, 0, sharcCapacity * sizeof(SharcPackedData));
    }
    
    vec3f camPos = camera.getFrom();
    float effectiveSceneScale = (sharcSceneScale > 0.0f) ? sharcSceneScale : this->sceneScale;
    
    SharcParameters sharcParams;
    sharcParams.gridParameters.cameraPosition = make_float3(camPos.x, camPos.y, camPos.z);
    sharcParams.gridParameters.logarithmBase = this->sharcLogBase;
    sharcParams.gridParameters.sceneScale = effectiveSceneScale;
    sharcParams.gridParameters.levelBias = (float)this->sharcLevelBias;
    
    sharcParams.hashMapData.capacity = sharcCapacity;
    sharcParams.hashMapData.hashEntriesBuffer = d_sharcHashEntries;
    sharcParams.hashMapData.lockBuffer = d_sharcLockBuffer;
    
    sharcParams.radianceScale = sharcRadianceScale;
    sharcParams.enableAntiFireflyFilter = sharcAntiFirefly;
    sharcParams.useHairTangent = sharcUseHairTangent;
    sharcParams.hairQueryMode = sharcHairQueryMode;
    sharcParams.hairTangentThreshold = sharcHairTangentThreshold;
    sharcParams.hairParallelScale = sharcHairParallelScale;
    sharcParams.hairPerpScale = sharcHairPerpScale;
    sharcParams.accumulationBuffer = d_sharcAccumBuffer;
    sharcParams.resolvedBuffer = d_sharcResolvedBuffer;
    sharcParams.hairMetaBuffer = d_sharcHairMetaBuffer;
    
    SharcResolveParameters resolveParams;
    resolveParams.cameraPositionPrev = sharcPrevCameraPos;
    resolveParams.accumulationFrameNum = sharcAccumFrameNum;
    resolveParams.staleFrameNumMax = sharcStaleFrameMax;
    resolveParams.enableAntiFireflyFilter = sharcAntiFirefly;
    resolveParams.blendAdjacentLevels = sharcBlendAdjacentLevels;
    
    const int blockSize = 256;
    const int numBlocks = (sharcCapacity + blockSize - 1) / blockSize;
    
    SharcResolveKernel<<<numBlocks, blockSize>>>(sharcParams, resolveParams);
    // No explicit sync needed: resolve kernel runs on default stream,
    // subsequent owlLaunch2D is serialized on the same stream.
    
    sharcPrevCameraPos = make_float3(camPos.x, camPos.y, camPos.z);
}

// ====================================================
// GPU Profiler
// ====================================================
void RenderWindowPT::initProfiler()
{
    if (profilerInitialized) return;
    cudaEventCreate(&evFrameStart);
    cudaEventCreate(&evFrameEnd);
    cudaEventCreate(&evUpdateStart);
    cudaEventCreate(&evUpdateEnd);
    cudaEventCreate(&evResolveStart);
    cudaEventCreate(&evResolveEnd);
    cudaEventCreate(&evRenderStart);
    cudaEventCreate(&evRenderEnd);
    profilerInitialized = true;
}

void RenderWindowPT::destroyProfiler()
{
    if (!profilerInitialized) return;
    cudaEventDestroy(evFrameStart);
    cudaEventDestroy(evFrameEnd);
    cudaEventDestroy(evUpdateStart);
    cudaEventDestroy(evUpdateEnd);
    cudaEventDestroy(evResolveStart);
    cudaEventDestroy(evResolveEnd);
    cudaEventDestroy(evRenderStart);
    cudaEventDestroy(evRenderEnd);
    profilerInitialized = false;
}

// ====================================================
// Custom functions
// ====================================================
vec3f RenderWindowPT::getSceneCg()
{
    return (this->minBound + this->maxBound) / 2.f;
}

void RenderWindowPT::initialize()
{
    vec2i frameRes(this->getWindowSize().x, this->getWindowSize().y);
    int frameSize = frameRes.x * frameRes.y;

    // Accumulation buffer stores the non-averaged radiance in float32
    // Float buffer stores the averaged radiance in float32
    // Frame buffer (defined in OWL::Viewer) stores averaged radiance in uint8
    cudaMalloc(&this->accumBuffer, frameSize * sizeof(float4));
    cudaMalloc(&this->averageBuffer, frameSize * sizeof(float4));

    // ====================================================
    // Environment lights setup
    // ====================================================
    if (this->currentScene.hasEnvLight) {

        int resx = this->currentScene.env.resolution.x;
        int resy = this->currentScene.env.resolution.y;
        vec2i resolution(resx, resy);
        float* envMap = this->currentScene.env.pixel_float;

        float* cPdf;
        float* cCdf;
        float* mPdf;
        float* mCdf;
        generateEnvSamplingTables(resolution, envMap, &cPdf, &cCdf, &mPdf, &mCdf);

        this->envTextureBuffer = owlTexture2DCreate(context,
            OWL_TEXEL_FORMAT_RGBA32F,
            resx, resy,
            envMap,
            OWL_TEXTURE_LINEAR,
            OWL_TEXTURE_CLAMP);
        this->envWidth = resx;
        this->envHeight = resy;

        this->conditionalPdf = owlTexture2DCreate(context,
            OWL_TEXEL_FORMAT_R32F,
            resx + 1, resy,
            cPdf,
            OWL_TEXTURE_NEAREST,
            OWL_TEXTURE_CLAMP);

        this->conditionalCdf = owlTexture2DCreate(context,
            OWL_TEXEL_FORMAT_R32F,
            resx + 1, resy,
            cCdf,
            OWL_TEXTURE_NEAREST,
            OWL_TEXTURE_CLAMP);

        this->marginalPdf = owlTexture2DCreate(context,
            OWL_TEXEL_FORMAT_R32F,
            resy + 1, 1.f,
            mPdf,
            OWL_TEXTURE_NEAREST,
            OWL_TEXTURE_CLAMP);

        this->marginalCdf = owlTexture2DCreate(context,
            OWL_TEXEL_FORMAT_R32F,
            resy + 1, 1.f,
            mCdf,
            OWL_TEXTURE_NEAREST,
            OWL_TEXTURE_CLAMP);

        free(cPdf); free(cCdf); free(mPdf); free(mCdf);
    }

    // ====================================================
    // Directional lights setup
    // ====================================================
    this->num_dlights = this->currentScene.dLightFrom.size();
    for (int i = 0; i < num_dlights; i++) {
        DirectionalLight dLight;
        dLight.from = this->currentScene.dLightFrom[i];
        dLight.emit = this->currentScene.dLightEmit[i];

        this->dLightList.push_back(dLight);
    }

    this->dLightsBuffer = owlDeviceBufferCreate(context,
        OWL_USER_TYPE(DirectionalLight), this->dLightList.size(), this->dLightList.data());

    // ====================================================
    // Vector of bottom level accel. structures (BLAS) 
    // This is then used to build the top level accel. structure (TLAS)
    // ====================================================
    std::vector<OWLGroup> blasList;

    // ====================================================
    // Surface geometry and data setup
    // ====================================================
    if (this->currentScene.has_surface) {
        auto surface = this->currentScene.surface;
        for (auto mesh : surface->meshes) {
            OWLVarDecl triangleGeomVars[] = {
                {"vertex", OWL_BUFPTR, OWL_OFFSETOF(TriangleMeshData, vertex)},
                {"normal", OWL_BUFPTR, OWL_OFFSETOF(TriangleMeshData, normal)},
                {"index", OWL_BUFPTR, OWL_OFFSETOF(TriangleMeshData, index)},
                {"texCoord", OWL_BUFPTR, OWL_OFFSETOF(TriangleMeshData, texCoord)},

                {"diffuse", OWL_FLOAT3, OWL_OFFSETOF(TriangleMeshData, diffuse)},
                {"diffuse_texture", OWL_TEXTURE, OWL_OFFSETOF(TriangleMeshData, diffuse_texture)},
                {"hasDiffuseTexture", OWL_BOOL, OWL_OFFSETOF(TriangleMeshData, hasDiffuseTexture)},

                {"alpha", OWL_FLOAT, OWL_OFFSETOF(TriangleMeshData, alpha)},
                {"alpha_texture", OWL_TEXTURE, OWL_OFFSETOF(TriangleMeshData, alpha_texture)},
                {"hasAlphaTexture", OWL_BOOL, OWL_OFFSETOF(TriangleMeshData, hasAlphaTexture)},

                {nullptr}
            };

            OWLGeomType triangleGeomType = owlGeomTypeCreate(context,
                OWL_GEOM_TRIANGLES,
                sizeof(TriangleMeshData),
                triangleGeomVars,
                -1);

            owlGeomTypeSetClosestHit(triangleGeomType, RADIANCE_RAY_TYPE, module, "triangleMeshCH");
            owlGeomTypeSetAnyHit(triangleGeomType, SHADOW_RAY_TYPE, module, "triangleMeshAHShadow");
            owlGeomTypeSetAnyHit(triangleGeomType, MULTISCATTER_RAY_TYPE, module, "triangleMeshAHMultiScatter");

            OWLGeom triangleGeom = owlGeomCreate(context, triangleGeomType);

            OWLBuffer vertexBuffer = owlDeviceBufferCreate(context, OWL_FLOAT3, mesh->vertex.size(), mesh->vertex.data());
            OWLBuffer vertexBuffer2 = owlDeviceBufferCreate(context, OWL_FLOAT3, mesh->vertex.size(), mesh->vertex.data());
            OWLBuffer normalBuffer = owlDeviceBufferCreate(context, OWL_FLOAT3, mesh->normal.size(), mesh->normal.data());
            OWLBuffer indexBuffer = owlDeviceBufferCreate(context, OWL_INT3, mesh->index.size(), mesh->index.data());
            OWLBuffer texCoordBuffer = owlDeviceBufferCreate(context, OWL_FLOAT2, mesh->texcoord.size(), mesh->texcoord.data());

            // Create CUDA buffers and upload them for diffuse and alpha textures
            if (mesh->diffuseTextureID != -1) {
                Texture* diffuseTexture = surface->textures[mesh->diffuseTextureID];
                OWLTexture diffuseTextureBuffer = owlTexture2DCreate(context,
                    OWL_TEXEL_FORMAT_RGBA8,
                    diffuseTexture->resolution.x,
                    diffuseTexture->resolution.y,
                    diffuseTexture->pixel,
                    OWL_TEXTURE_NEAREST,
                    OWL_TEXTURE_MIRROR);
                owlGeomSetTexture(triangleGeom, "diffuse_texture", diffuseTextureBuffer);
                owlGeomSet1b(triangleGeom, "hasDiffuseTexture", true);
            }
            else {
                owlGeomSet3f(triangleGeom, "diffuse", owl3f{ mesh->diffuse.x, mesh->diffuse.y, mesh->diffuse.z });
                owlGeomSet1b(triangleGeom, "hasDiffuseTexture", false);
            }

            if (mesh->alphaTextureID != -1) {
                Texture* alphaTexture = surface->textures[mesh->alphaTextureID];
                OWLTexture alphaTextureBuffer = owlTexture2DCreate(context,
                    OWL_TEXEL_FORMAT_RGBA8,
                    alphaTexture->resolution.x,
                    alphaTexture->resolution.y,
                    alphaTexture->pixel,
                    OWL_TEXTURE_NEAREST,
                    OWL_TEXTURE_MIRROR);
                owlGeomSetTexture(triangleGeom, "alpha_texture", alphaTextureBuffer);
                owlGeomSet1b(triangleGeom, "hasAlphaTexture", true);
            }
            else {
                owlGeomSet1f(triangleGeom, "alpha", mesh->alpha);
                owlGeomSet1b(triangleGeom, "hasAlphaTexture", false);
            }

            owlTrianglesSetVertices(triangleGeom, vertexBuffer,
                mesh->vertex.size(), sizeof(vec3f), 0);
            owlTrianglesSetIndices(triangleGeom, indexBuffer,
                mesh->index.size(), sizeof(vec3i), 0);

            owlGeomSetBuffer(triangleGeom, "vertex", vertexBuffer2);
            owlGeomSetBuffer(triangleGeom, "normal", normalBuffer);
            owlGeomSetBuffer(triangleGeom, "index", indexBuffer);
            owlGeomSetBuffer(triangleGeom, "texCoord", texCoordBuffer);

            OWLGroup triangleGroup = owlTrianglesGeomGroupCreate(context, 1, &triangleGeom);
            owlGroupBuildAccel(triangleGroup);

            this->minBound = min(this->minBound, mesh->minBound);
            this->maxBound = max(this->maxBound, mesh->maxBound);

            // Add to a list, which is later used to build the IAS (TLAS)
            blasList.push_back(triangleGroup);
        }

        this->sceneScale = length(this->maxBound - this->minBound);
    }


    // ====================================================
    // Hair setup (geometry and BSDF)
    // ====================================================
    OWLVarDecl hairGeomVars[] = {
        {nullptr}
    };

    OWLGeomType hairGeomType = owlGeomTypeCreate(context,
        OWL_GEOMETRY_CURVES,
        sizeof(HairData),
        hairGeomVars,
        -1);

    owlGeomTypeSetClosestHit(hairGeomType, RADIANCE_RAY_TYPE, module, "hairCH");
    owlGeomTypeSetAnyHit(hairGeomType, SHADOW_RAY_TYPE, module, "hairAHShadow");
    owlGeomTypeSetAnyHit(hairGeomType, MULTISCATTER_RAY_TYPE, module, "hairAHMultiScatter");
    owlCurvesSetDegree(hairGeomType, 3, true);

    // ====================================================
    // Data setup (hair strands from .hair file)
    // ====================================================

    this->hairGeom = owlGeomCreate(context, hairGeomType);

    if(this->currentScene.has_hair) {
        this->minBound = min(this->currentScene.hairModel.minBound, minBound);
        this->maxBound = max(this->currentScene.hairModel.maxBound, maxBound);
        this->sceneScale = max(this->currentScene.hairModel.scale, this->sceneScale);

        OWLBuffer cpBuffer = owlDeviceBufferCreate(context, OWL_FLOAT3, 
            this->currentScene.hairModel.controlPoints.size(), this->currentScene.hairModel.controlPoints.data());
        OWLBuffer widthsBuffer = owlDeviceBufferCreate(context, OWL_FLOAT, 
            this->currentScene.hairModel.widths.size(), this->currentScene.hairModel.widths.data());
        OWLBuffer indicesBuffer = owlDeviceBufferCreate(context, OWL_INT, 
            this->currentScene.hairModel.segmentIndices.size(), this->currentScene.hairModel.segmentIndices.data());

        owlCurvesSetControlPoints(this->hairGeom, this->currentScene.hairModel.controlPoints.size(), cpBuffer, widthsBuffer);
        owlCurvesSetSegmentIndices(this->hairGeom, this->currentScene.hairModel.segmentIndices.size(), indicesBuffer);

        OWLGroup hairGroup = owlCurvesGeomGroupCreate(context, 1, &this->hairGeom);
        owlGroupBuildAccel(hairGroup);

        blasList.push_back(hairGroup);
    }

    // ====================================================
    // Build the TLAS (IAS)
    // ====================================================
    world = owlInstanceGroupCreate(context, blasList.size(), blasList.data());
    owlGroupBuildAccel(world);

    // ====================================================
    // Launch Parameters setup
    // ====================================================
    OWLVarDecl launchParamsDecl[] = {
        // Environment light parametsrs
        {"envRotPhi", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, envRotPhi)},
        {"envScale", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, envScale)},
        {"hasEnvLight", OWL_BOOL, OWL_OFFSETOF(LaunchParams, hasEnvLight)},
        {"env", OWL_TEXTURE, OWL_OFFSETOF(LaunchParams, env)},
        {"envWidth", OWL_INT, OWL_OFFSETOF(LaunchParams, envWidth)},
        {"envHeight", OWL_INT, OWL_OFFSETOF(LaunchParams, envHeight)},
        // Environment light PDF & CDF
        {"conditionalPdf", OWL_TEXTURE, OWL_OFFSETOF(LaunchParams, conditionalPdf)},
        {"conditionalCdf", OWL_TEXTURE, OWL_OFFSETOF(LaunchParams, conditionalCdf)},
        {"marginalPdf", OWL_TEXTURE, OWL_OFFSETOF(LaunchParams, marginalPdf)},
        {"marginalCdf", OWL_TEXTURE, OWL_OFFSETOF(LaunchParams, marginalCdf)},
        // Directional light parameters
        {"hasDirectionalLights", OWL_BOOL, OWL_OFFSETOF(LaunchParams, hasDirectionalLights)},
        {"dLights", OWL_BUFPTR, OWL_OFFSETOF(LaunchParams, dLights)},
        {"num_dlights", OWL_INT, OWL_OFFSETOF(LaunchParams, num_dlights)},
        // Global light parameters
        {"num_total_lights", OWL_INT, OWL_OFFSETOF(LaunchParams, num_total_lights)},
        // Integrator parameters
        {"pathV1", OWL_INT, OWL_OFFSETOF(LaunchParams, pathV1)},
        {"pathV2", OWL_INT, OWL_OFFSETOF(LaunchParams, pathV2)},
        {"MIS", OWL_BOOL, OWL_OFFSETOF(LaunchParams, MIS)},
        {"envPdfSampling", OWL_BOOL, OWL_OFFSETOF(LaunchParams, envPdfSampling)},
        // Scene properties
        {"maxBound", OWL_FLOAT3, OWL_OFFSETOF(LaunchParams, maxBound)},
        {"minBound", OWL_FLOAT3, OWL_OFFSETOF(LaunchParams, minBound)},
        {"sceneScale", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, sceneScale)},
        // Global Hair parameters
        {"sig_a", OWL_FLOAT3, OWL_OFFSETOF(LaunchParams, hairData.sig_a)},
        {"beta_m", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, hairData.beta_m)},
        {"beta_n", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, hairData.beta_n)},
        {"alpha", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, hairData.alpha)},
        {"R_G", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, hairData.R_G)},
        {"TT_G", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, hairData.TT_G)},
        {"TRT_G", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, hairData.TRT_G)},
        {"TRRT_G", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, hairData.TRRT_G)},
        // All other parameters
        {"accumBuffer", OWL_RAW_POINTER, OWL_OFFSETOF(LaunchParams, accumBuffer)},
        {"averageBuffer", OWL_RAW_POINTER, OWL_OFFSETOF(LaunchParams, averageBuffer)},
        {"accumId", OWL_INT, OWL_OFFSETOF(LaunchParams, accumId)},
        {"oneSppRealtime", OWL_BOOL, OWL_OFFSETOF(LaunchParams, oneSppRealtime)},
        {"world", OWL_GROUP, OWL_OFFSETOF(LaunchParams, world)},
        {"camera.pos", OWL_FLOAT3, OWL_OFFSETOF(LaunchParams, camera.pos)},
        {"camera.dir_00", OWL_FLOAT3, OWL_OFFSETOF(LaunchParams, camera.dir_00)},
        {"camera.dir_du", OWL_FLOAT3, OWL_OFFSETOF(LaunchParams, camera.dir_du)},
        {"camera.dir_dv", OWL_FLOAT3, OWL_OFFSETOF(LaunchParams, camera.dir_dv)},
        // SHARC parameters
        {"sharcEnabled", OWL_BOOL, OWL_OFFSETOF(LaunchParams, sharcEnabled)},
        {"sharcUpdateEnabled", OWL_BOOL, OWL_OFFSETOF(LaunchParams, sharcUpdateEnabled)},
        {"sharcLaunchMode", OWL_INT, OWL_OFFSETOF(LaunchParams, sharcLaunchMode)},
        {"sharcDebugMode", OWL_INT, OWL_OFFSETOF(LaunchParams, sharcDebugMode)},
        {"sharcUpdateWidth", OWL_INT, OWL_OFFSETOF(LaunchParams, sharcUpdateWidth)},
        {"sharcUpdateHeight", OWL_INT, OWL_OFFSETOF(LaunchParams, sharcUpdateHeight)},
        {"sharcBounce", OWL_INT, OWL_OFFSETOF(LaunchParams, sharcBounce)},
        {"sharcMaxBounces", OWL_INT, OWL_OFFSETOF(LaunchParams, sharcMaxBounces)},
        {"sharcSampleThreshold", OWL_INT, OWL_OFFSETOF(LaunchParams, sharcSampleThreshold)},
        {"sharcCompareMode", OWL_INT, OWL_OFFSETOF(LaunchParams, sharcCompareMode)},
        {"sharcQueryScale", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, sharcQueryScale)},
        {"sharcHitDistanceScale", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, sharcHitDistanceScale)},
        {"sharcQueryJitter", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, sharcQueryJitter)},
        {"sharcShowBeta", OWL_INT, OWL_OFFSETOF(LaunchParams, sharcShowBeta)},
        {"sharcParams.gridParameters.cameraPosition", OWL_FLOAT3, OWL_OFFSETOF(LaunchParams, sharcParams.gridParameters.cameraPosition)},
        {"sharcParams.gridParameters.logarithmBase", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, sharcParams.gridParameters.logarithmBase)},
        {"sharcParams.gridParameters.sceneScale", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, sharcParams.gridParameters.sceneScale)},
        {"sharcParams.gridParameters.levelBias", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, sharcParams.gridParameters.levelBias)},
        {"sharcParams.hashMapData.capacity", OWL_UINT, OWL_OFFSETOF(LaunchParams, sharcParams.hashMapData.capacity)},
        {"sharcParams.hashMapData.hashEntriesBuffer", OWL_RAW_POINTER, OWL_OFFSETOF(LaunchParams, sharcParams.hashMapData.hashEntriesBuffer)},
        {"sharcParams.hashMapData.lockBuffer", OWL_RAW_POINTER, OWL_OFFSETOF(LaunchParams, sharcParams.hashMapData.lockBuffer)},
        {"sharcParams.radianceScale", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, sharcParams.radianceScale)},
        {"sharcParams.enableAntiFireflyFilter", OWL_BOOL, OWL_OFFSETOF(LaunchParams, sharcParams.enableAntiFireflyFilter)},
        {"sharcParams.disableEnvMiss", OWL_BOOL, OWL_OFFSETOF(LaunchParams, sharcParams.disableEnvMiss)},
        {"sharcParams.enableMaterialDemodulation", OWL_BOOL, OWL_OFFSETOF(LaunchParams, sharcParams.enableMaterialDemodulation)},
        {"sharcParams.useHairTangent", OWL_BOOL, OWL_OFFSETOF(LaunchParams, sharcParams.useHairTangent)},
        {"sharcParams.hairQueryMode", OWL_INT, OWL_OFFSETOF(LaunchParams, sharcParams.hairQueryMode)},
        {"sharcParams.hairTangentThreshold", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, sharcParams.hairTangentThreshold)},
        {"sharcParams.hairParallelScale", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, sharcParams.hairParallelScale)},
        {"sharcParams.hairPerpScale", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, sharcParams.hairPerpScale)},
        {"sharcDemodScale", OWL_FLOAT, OWL_OFFSETOF(LaunchParams, sharcDemodScale)},
        {"sharcParams.accumulationBuffer", OWL_RAW_POINTER, OWL_OFFSETOF(LaunchParams, sharcParams.accumulationBuffer)},
        {"sharcParams.resolvedBuffer", OWL_RAW_POINTER, OWL_OFFSETOF(LaunchParams, sharcParams.resolvedBuffer)},
        {"sharcParams.hairMetaBuffer", OWL_RAW_POINTER, OWL_OFFSETOF(LaunchParams, sharcParams.hairMetaBuffer)},
        {nullptr}
    };

    this->launchParams = owlParamsCreate(context, sizeof(LaunchParams), launchParamsDecl, -1);

    // Global hair parameters
    owlParamsSet3f(this->launchParams, "sig_a", owl3f{ this->currentScene.sig_a.x, this->currentScene.sig_a.y, this->currentScene.sig_a.z });
    owlParamsSet1f(this->launchParams, "beta_m", this->currentScene.beta_m);
    owlParamsSet1f(this->launchParams, "beta_n", this->currentScene.beta_n);
    owlParamsSet1f(this->launchParams, "alpha", this->currentScene.alpha);

    owlParamsSet1f(this->launchParams, "R_G", this->currentScene.R_G);
    owlParamsSet1f(this->launchParams, "TT_G", this->currentScene.TT_G);
    owlParamsSet1f(this->launchParams, "TRT_G", this->currentScene.TRT_G);
    owlParamsSet1f(this->launchParams, "TRRT_G", this->currentScene.TRRT_G);

    // Upload scene properties
    owlParamsSet3f(this->launchParams, "maxBound", owl3f{ this->maxBound.x, this->maxBound.y, this->maxBound.z });
    owlParamsSet3f(this->launchParams, "minBound", owl3f{ this->minBound.x, this->minBound.y, this->minBound.z });
    owlParamsSet1f(this->launchParams, "sceneScale", this->sceneScale);

    // Upload integrator parameters
    owlParamsSet1b(this->launchParams, "MIS", this->currentScene.MIS);
    owlParamsSet1i(this->launchParams, "pathV1", this->currentScene.path_v1);
    owlParamsSet1i(this->launchParams, "pathV2", this->currentScene.path_v2);
    owlParamsSet1i(this->launchParams, "accumId", this->accumId);
    owlParamsSet1b(this->launchParams, "oneSppRealtime", this->oneSppRealtime);
    owlParamsSet1ul(this->launchParams, "accumBuffer", (uint64_t)this->accumBuffer);
    owlParamsSet1ul(this->launchParams, "averageBuffer", (uint64_t)this->averageBuffer);



    int num_total_lights = 0;

    // Upload environment light
    owlParamsSet1f(this->launchParams, "envRotPhi", this->currentScene.envRotPhi);
    owlParamsSet1f(this->launchParams, "envScale", this->currentScene.envScale);
    owlParamsSet1i(this->launchParams, "envWidth", this->envWidth);
    owlParamsSet1i(this->launchParams, "envHeight", this->envHeight);
    owlParamsSet1b(this->launchParams, "envPdfSampling", this->currentScene.envPdfSampling);
    owlParamsSet1b(this->launchParams, "hasEnvLight", this->currentScene.hasEnvLight);
    if (this->currentScene.hasEnvLight) {
        owlParamsSetTexture(this->launchParams, "env", this->envTextureBuffer);

        owlParamsSetTexture(this->launchParams, "conditionalPdf", this->conditionalPdf);
        owlParamsSetTexture(this->launchParams, "conditionalCdf", this->conditionalCdf);
        owlParamsSetTexture(this->launchParams, "marginalPdf", this->marginalPdf);
        owlParamsSetTexture(this->launchParams, "marginalCdf", this->marginalCdf);

        num_total_lights += 1;
    }

    // Upload directional lights
    owlParamsSet1b(this->launchParams, "hasDirectionalLights", this->currentScene.hasDirectionalLights);
    owlParamsSetBuffer(this->launchParams, "dLights", dLightsBuffer);
    owlParamsSet1i(this->launchParams, "num_dlights", num_dlights);

    // Upload global light parameters
    num_total_lights += num_dlights;
    owlParamsSet1i(this->launchParams, "num_total_lights", num_total_lights);

    // ====================================================
    // Setup a generic miss program
    // ====================================================
    OWLVarDecl missProgVars[] = {
        {"const_color", OWL_FLOAT3, OWL_OFFSETOF(MissProgData, const_color)},
        {nullptr}
    };

    missProg = owlMissProgCreate(context, module, "miss", sizeof(MissProgData), missProgVars, -1);

    // Set a constant background color in the miss program (black for now)
    owlMissProgSet3f(missProg, "const_color", owl3f{ 0.f, 0.f, 0.f });

    // ====================================================
    // Setup a pin-hole camera ray-gen program
    // ====================================================
    OWLVarDecl rayGenVars[] = {
        {"frameBuffer", OWL_RAW_POINTER, OWL_OFFSETOF(RayGenData, frameBuffer)},
        {"frameBufferSize", OWL_INT2, OWL_OFFSETOF(RayGenData, frameBufferSize)},
        {nullptr}
    };

    // Create regular ray gen
    this->rayGen = owlRayGenCreate(context, module, "rayGenCam", sizeof(RayGenData), rayGenVars, -1);

    // Set the TLAS to be used
    owlParamsSetGroup(this->launchParams, "world", world);

    // ====================================================-
    
    // SHARC initialization (default disabled)
    // ====================================================
    // Set SHARC parameters (buffers are null by default - will be allocated when enabled)
    owlParamsSet1b(this->launchParams, "sharcEnabled", this->sharcEnabled);
    owlParamsSet1i(this->launchParams, "sharcCompareMode",this->sharcCompareMode);
    owlParamsSet1i(this->launchParams, "sharcShowBeta", this->sharcShowBeta);
    owlParamsSet1b(this->launchParams, "sharcUpdateEnabled", this->sharcUpdateEnabled);
    owlParamsSet1i(this->launchParams, "sharcBounce", this->sharcBounce);
    owlParamsSet1i(this->launchParams, "sharcDebugMode", this->sharcDebugMode);
    
    vec3f camPos = camera.getFrom();
    owlParamsSet3f(this->launchParams, "sharcParams.gridParameters.cameraPosition", 
                   owl3f{ camPos.x, camPos.y, camPos.z });
    owlParamsSet1f(this->launchParams, "sharcParams.gridParameters.logarithmBase", SHARC_GRID_LOGARITHM_BASE);
    owlParamsSet1f(this->launchParams, "sharcParams.gridParameters.sceneScale", this->sceneScale);
    owlParamsSet1f(this->launchParams, "sharcParams.gridParameters.levelBias", SHARC_GRID_LEVEL_BIAS);
    
    owlParamsSet1ui(this->launchParams, "sharcParams.hashMapData.capacity", sharcCapacity);
    owlParamsSet1ul(this->launchParams, "sharcParams.hashMapData.hashEntriesBuffer", 0);  // null initially
    owlParamsSet1ul(this->launchParams, "sharcParams.hashMapData.lockBuffer", 0);
    owlParamsSet1f(this->launchParams, "sharcParams.radianceScale", sharcRadianceScale);
    owlParamsSet1b(this->launchParams, "sharcParams.enableAntiFireflyFilter", sharcAntiFirefly);
    owlParamsSet1b(this->launchParams, "sharcParams.disableEnvMiss", sharcDisableEnvMiss);
    owlParamsSet1b(this->launchParams, "sharcParams.enableMaterialDemodulation", sharcMaterialDemod);
    owlParamsSet1b(this->launchParams, "sharcParams.useHairTangent", sharcUseHairTangent);
    owlParamsSet1i(this->launchParams, "sharcParams.hairQueryMode", sharcHairQueryMode);
    owlParamsSet1f(this->launchParams, "sharcParams.hairTangentThreshold", sharcHairTangentThreshold);
    owlParamsSet1f(this->launchParams, "sharcParams.hairParallelScale", sharcHairParallelScale);
    owlParamsSet1f(this->launchParams, "sharcParams.hairPerpScale", sharcHairPerpScale);
    owlParamsSet1f(this->launchParams, "sharcDemodScale", sharcDemodScale);
    owlParamsSet1ul(this->launchParams, "sharcParams.accumulationBuffer", 0);
    owlParamsSet1ul(this->launchParams, "sharcParams.resolvedBuffer", 0);
    owlParamsSet1ul(this->launchParams, "sharcParams.hairMetaBuffer", 0);

    if (sharcEnabled && !sharcInitialized) {
        initializeSharc();
    }

    // ====================================================
    // Finally, build the programs, pipeline and SBT
    // ====================================================
    owlBuildPrograms(context);
    owlBuildPipeline(context);
    owlBuildSBT(context);

    this->enableInspectMode(owl::box3f(this->minBound, this->maxBound));
    this->setWorldScale(this->sceneScale);

    // Initialize saveIndex by scanning for existing files to avoid overwriting
    {
        std::string basePath = this->currentScene.renderOutput;
        size_t dotPos = basePath.rfind('.');
        std::string stem = (dotPos != std::string::npos) ? basePath.substr(0, dotPos) : basePath;
        
        this->saveIndex = 0;
        char testName[256];
        struct stat st;
        while (this->saveIndex < 10000) {
            // Check both .png and .exr variants
            snprintf(testName, sizeof(testName), "%s%03d.png", stem.c_str(), this->saveIndex);
            bool pngExists = (stat(testName, &st) == 0);
            snprintf(testName, sizeof(testName), "%s%03d.exr", stem.c_str(), this->saveIndex);
            bool exrExists = (stat(testName, &st) == 0);
            if (!pngExists && !exrExists)
                break;
            this->saveIndex++;
        }
        if (this->saveIndex > 0)
            printf("Save index starts at %03d (found existing files)\n", this->saveIndex);
    }
}

long long RenderWindowPT::render()
{
    auto start = std::chrono::high_resolution_clock::now();

    // Initialize profiler on first use
    if (profilerEnabled && !profilerInitialized) {
        initProfiler();
    }

    if (sbtDirty) {
        owlBuildSBT(context);
        sbtDirty = false;
    }

    if (taskTimingEnabled && !taskTimingInitialized) {
        cudaEventCreate(&taskEvRenderStart);
        cudaEventCreate(&taskEvRenderEnd);
        taskTimingInitialized = true;
    }

    const bool runOneSppRealtime = this->oneSppRealtime;
    if (runOneSppRealtime || this->accumId < this->currentScene.spp || this->progressive) {
        // Update SHARC parameters each frame
        owlParamsSet1b(this->launchParams, "sharcEnabled", this->sharcEnabled);
        owlParamsSet1b(this->launchParams, "sharcUpdateEnabled", this->sharcUpdateEnabled);
        owlParamsSet1i(this->launchParams, "sharcDebugMode", this->sharcDebugMode);
        owlParamsSet1i(this->launchParams, "sharcBounce", this->sharcBounce);
        owlParamsSet1i(this->launchParams, "sharcMaxBounces", this->sharcMaxBounces);
        owlParamsSet1i(this->launchParams, "sharcSampleThreshold", this->sharcSampleThreshold);
        owlParamsSet1i(this->launchParams, "sharcCompareMode",this->sharcCompareMode);
        owlParamsSet1f(this->launchParams, "sharcQueryScale", this->sharcQueryScale);
        owlParamsSet1f(this->launchParams, "sharcHitDistanceScale", this->sharcHitDistanceScale);
        owlParamsSet1f(this->launchParams, "sharcQueryJitter", this->sharcQueryJitter);
        owlParamsSet1i(this->launchParams, "sharcShowBeta", this->sharcShowBeta);
        owlParamsSet1i(this->launchParams, "accumId", this->accumId);
        owlParamsSet1b(this->launchParams, "oneSppRealtime", runOneSppRealtime);
        
        // Update grid parameters
        if (sharcEnabled) {
            vec3f camPos = camera.getFrom();
            owlParamsSet3f(this->launchParams, "sharcParams.gridParameters.cameraPosition", 
                           owl3f{ camPos.x, camPos.y, camPos.z });
            float effectiveSceneScale = (sharcSceneScale > 0.0f) ? sharcSceneScale : this->sceneScale;
            owlParamsSet1f(this->launchParams, "sharcParams.gridParameters.sceneScale", effectiveSceneScale);
            owlParamsSet1f(this->launchParams, "sharcParams.gridParameters.logarithmBase", this->sharcLogBase);
            owlParamsSet1f(this->launchParams, "sharcParams.gridParameters.levelBias", (float)this->sharcLevelBias);
            owlParamsSet1f(this->launchParams, "sharcParams.radianceScale", this->sharcRadianceScale);
            owlParamsSet1b(this->launchParams, "sharcParams.disableEnvMiss", this->sharcDisableEnvMiss);
            owlParamsSet1b(this->launchParams, "sharcParams.enableMaterialDemodulation", this->sharcMaterialDemod);
            owlParamsSet1b(this->launchParams, "sharcParams.useHairTangent", this->sharcUseHairTangent);
            owlParamsSet1i(this->launchParams, "sharcParams.hairQueryMode", this->sharcHairQueryMode);
            owlParamsSet1f(this->launchParams, "sharcParams.hairTangentThreshold", this->sharcHairTangentThreshold);
            owlParamsSet1f(this->launchParams, "sharcParams.hairParallelScale", this->sharcHairParallelScale);
            owlParamsSet1f(this->launchParams, "sharcParams.hairPerpScale", this->sharcHairPerpScale);
            owlParamsSet1f(this->launchParams, "sharcDemodScale", this->sharcDemodScale);
        }
        
        // Calculate update resolution
        int updateScale = max(1, this->sharcUpdateResScale);
        int updateWidth = max(1, this->fbSize.x / updateScale);
        int updateHeight = max(1, this->fbSize.y / updateScale);
        owlParamsSet1i(this->launchParams, "sharcUpdateWidth", updateWidth);
        owlParamsSet1i(this->launchParams, "sharcUpdateHeight", updateHeight);
        
        const bool prof = profilerEnabled && profilerInitialized;
        
        if (prof) cudaEventRecord(evFrameStart);
        
        // ========== PASS 1: SHARC Update (low resolution) ==========
        if (sharcEnabled && sharcUpdateEnabled && sharcInitialized) {
            if (prof) cudaEventRecord(evUpdateStart);
            
            owlParamsSet1i(this->launchParams, "sharcLaunchMode", 1);  // Update mode
            owlLaunch2D(this->rayGen, updateWidth, updateHeight, this->launchParams);
            
            if (prof) cudaEventRecord(evUpdateEnd);
            
            // ========== PASS 1.5: SHARC Resolve ==========
            if (prof) cudaEventRecord(evResolveStart);
            resolveSharc();
            if (prof) cudaEventRecord(evResolveEnd);
        }
        
        // ========== PASS 2: Render or Debug (full resolution) ==========
        if (prof) cudaEventRecord(evRenderStart);
        
        if (sharcDebugMode > 0 && sharcEnabled) {
            owlParamsSet1i(this->launchParams, "sharcLaunchMode", 2);  // Debug mode
        } else {
            owlParamsSet1i(this->launchParams, "sharcLaunchMode", 0);  // Render mode
        }
        if (taskTimingEnabled && taskTimingInitialized)
            cudaEventRecord(taskEvRenderStart);
        owlLaunch2D(this->rayGen, this->fbSize.x, this->fbSize.y, this->launchParams);
        if (taskTimingEnabled && taskTimingInitialized) {
            cudaEventRecord(taskEvRenderEnd);
            cudaEventSynchronize(taskEvRenderEnd);
            float renderMs = 0.f;
            cudaEventElapsedTime(&renderMs, taskEvRenderStart, taskEvRenderEnd);
            taskGpuRenderMsLast = renderMs;
            taskGpuRenderMsTotal += (double)renderMs;
        }
        
        if (prof) cudaEventRecord(evRenderEnd);
        if (prof) cudaEventRecord(evFrameEnd);
        
        // Collect profiler results (from previous frame's events)
        if (prof) {
            cudaEventSynchronize(evFrameEnd);
            
            float totalMs = 0.f;
            cudaEventElapsedTime(&totalMs, evFrameStart, evFrameEnd);
            profTotalMs += PROF_SMOOTH_FACTOR * (totalMs - profTotalMs);
            
            float renderMs = 0.f;
            cudaEventElapsedTime(&renderMs, evRenderStart, evRenderEnd);
            profRenderMs += PROF_SMOOTH_FACTOR * (renderMs - profRenderMs);
            
            if (sharcEnabled && sharcUpdateEnabled && sharcInitialized) {
                float updateMs = 0.f;
                cudaEventElapsedTime(&updateMs, evUpdateStart, evUpdateEnd);
                profUpdateMs += PROF_SMOOTH_FACTOR * (updateMs - profUpdateMs);
                
                float resolveMs = 0.f;
                cudaEventElapsedTime(&resolveMs, evResolveStart, evResolveEnd);
                profResolveMs += PROF_SMOOTH_FACTOR * (resolveMs - profResolveMs);
            } else {
                profUpdateMs = 0.f;
                profResolveMs = 0.f;
            }
        }

        accumId++;
    }

    auto finish = std::chrono::high_resolution_clock::now();

    return std::chrono::duration_cast<std::chrono::microseconds>(finish - start).count();
}

void RenderWindowPT::resize(const vec2i& newSize)
{
    if (this->getWindowSize().x == 0 && this->getWindowSize().y == 0) {
        OWLViewer::resize(newSize);

        this->initialize();
        this->cameraChanged();
    }
}

void RenderWindowPT::drawUI()
{
    ImGui_ImplOpenGL2_NewFrame();
    ImGui_ImplGlfw_NewFrame();
    ImGui::NewFrame();

    {
        ImGui::Begin("Controls", NULL, ImGuiWindowFlags_AlwaysAutoResize | ImGuiWindowFlags_NoMove | ImGuiWindowFlags_NoResize);
        ImGui::Text("%.3f ms/frame (%.1f FPS)", 1000.0f / ImGui::GetIO().Framerate, ImGui::GetIO().Framerate);

        bool prog = this->progressive;
        ImGui::Checkbox("Progressive", &prog);
        if (prog != this->progressive) {
            this->progressive = prog;
            // Keep sampling modes mutually exclusive:
            // enabling progressive disables 1-SPP realtime mode.
            if (this->progressive && this->oneSppRealtime) {
                this->oneSppRealtime = false;
            }
            this->cameraChanged();
        }

        bool oneSppRt = this->oneSppRealtime;
        ImGui::Checkbox("1 SPP Realtime", &oneSppRt);
        ImGui::SetItemTooltip("Render every frame with accumId=0 (no temporal accumulation).");
        if (oneSppRt != this->oneSppRealtime) {
            this->oneSppRealtime = oneSppRt;
            // Keep sampling modes mutually exclusive:
            // enabling 1-SPP realtime disables progressive accumulation.
            if (this->oneSppRealtime && this->progressive) {
                this->progressive = false;
            }
            this->cameraChanged();
        }

        if (!this->progressive && !this->oneSppRealtime) {
            int currentSpp = this->currentScene.spp;
            ImGui::SliderInt("SPP", &currentSpp, 1, 100);
            if (currentSpp != this->currentScene.spp) {
                this->currentScene.spp = currentSpp;
                this->cameraChanged();
            }
        }

        // Sampling controls
        if (ImGui::CollapsingHeader("Sampling")) {
            bool currentMIS = this->currentScene.MIS;
            ImGui::Checkbox("MIS", &currentMIS);
            if (currentMIS != this->currentScene.MIS) {
                this->currentScene.MIS = currentMIS;
                owlParamsSet1b(this->launchParams, "MIS", currentMIS);
                this->cameraChanged();
            }

            bool currentEnvSampling = this->currentScene.envPdfSampling;
            ImGui::Checkbox("Env. PDF Samp.", &currentEnvSampling);
            if (currentEnvSampling != this->currentScene.envPdfSampling) {
                this->currentScene.envPdfSampling = currentEnvSampling;
                owlParamsSet1b(this->launchParams, "envPdfSampling", currentEnvSampling);
                this->cameraChanged();
            }
        }

        // All hair controls
        // All hair controls
        if (ImGui::CollapsingHeader("Hair")) {
            float cur_sig_a[3] = { this->currentScene.sig_a.x, this->currentScene.sig_a.y,
                                                this->currentScene.sig_a.z };
            ImGui::SliderFloat3("Sigma", cur_sig_a, 0.f, 3.f);
            if (this->currentScene.sig_a.x != cur_sig_a[0] ||
                this->currentScene.sig_a.y != cur_sig_a[1] ||
                this->currentScene.sig_a.z != cur_sig_a[2]) {
                this->currentScene.sig_a = vec3f(cur_sig_a[0], cur_sig_a[1], cur_sig_a[2]);
                owlParamsSet3f(this->launchParams, "sig_a", owl3f{ cur_sig_a[0], cur_sig_a[1],
                                                                cur_sig_a[2] });
                this->cameraChanged();
            }

            float cur_beta_m = this->currentScene.beta_m;
            ImGui::SliderFloat("beta_m", &cur_beta_m, 0.f, 1.f);
            if (cur_beta_m != this->currentScene.beta_m) {
                this->currentScene.beta_m = cur_beta_m;
                owlParamsSet1f(this->launchParams, "beta_m", cur_beta_m);
                this->cameraChanged();
            }

            float cur_beta_n = this->currentScene.beta_n;
            ImGui::SliderFloat("beta_n", &cur_beta_n, 0.f, 1.f);
            if (cur_beta_n != this->currentScene.beta_n) {
                this->currentScene.beta_n = cur_beta_n;
                owlParamsSet1f(this->launchParams, "beta_n", cur_beta_n);
                this->cameraChanged();
            }

            float cur_alpha = this->currentScene.alpha;
            ImGui::SliderFloat("alpha", &cur_alpha, 0.f, 20.f * 3.14159f / 180.f);
            if (cur_alpha != this->currentScene.alpha) {
                this->currentScene.alpha = cur_alpha;
                owlParamsSet1f(this->launchParams, "alpha", cur_alpha);
                this->cameraChanged();
            }

            float rg = this->currentScene.R_G;
            ImGui::SliderFloat("R Gain", &rg, 0.f, 1.f);
            if (rg != this->currentScene.R_G) {
                this->currentScene.R_G = rg;
                owlParamsSet1f(this->launchParams, "R_G", rg);
                this->cameraChanged();
            }

            float ttg = this->currentScene.TT_G;
            ImGui::SliderFloat("TT Gain", &ttg, 0.f, 1.f);
            if (ttg != this->currentScene.TT_G) {
                this->currentScene.TT_G = ttg;
                owlParamsSet1f(this->launchParams, "TT_G", ttg);
                this->cameraChanged();
            }

            float trtg = this->currentScene.TRT_G;
            ImGui::SliderFloat("TRT Gain", &trtg, 0.f, 1.f);
            if (trtg != this->currentScene.TRT_G) {
                this->currentScene.TRT_G = trtg;
                owlParamsSet1f(this->launchParams, "TRT_G", trtg);
                this->cameraChanged();
            }

            float trrtg = this->currentScene.TRRT_G;
            ImGui::SliderFloat("TRRT Gain", &trrtg, 0.f, 1.f);
            if (trrtg != this->currentScene.TRRT_G) {
                this->currentScene.TRRT_G = trrtg;
                owlParamsSet1f(this->launchParams, "TRRT_G", trrtg);
                this->cameraChanged();
            }
        }

        // SHARC Cache controls
        if (ImGui::CollapsingHeader("SHARC Cache")) {
            bool enabled = this->sharcEnabled;
            ImGui::Checkbox("Enable SHARC", &enabled);
            if (enabled && !sharcInitialized) {
                initializeSharc();
            }
            if (enabled != this->sharcEnabled) {
                this->sharcEnabled = enabled;
                this->cameraChanged();
            }
            
            if (this->sharcEnabled) {
                // Main controls
                bool update = this->sharcUpdateEnabled;
                ImGui::Checkbox("Enable Update", &update);
                if (update != this->sharcUpdateEnabled) {
                    this->sharcUpdateEnabled = update;
                }
                
                const char* debugModes[] = { "Off", "Sample Count", "Radiance", "Voxel ID", "Voxel Size", "Grid Level", "Raw Accum", "HitDist/Voxel" };
                int debugMode = this->sharcDebugMode;
                ImGui::Combo("Debug Mode", &debugMode, debugModes, 8);
                if (debugMode != this->sharcDebugMode) {
                    this->sharcDebugMode = debugMode;
                    this->cameraChanged();
                }
                
                ImGui::Separator();
                
                // Core settings
                if (ImGui::TreeNode("Settings")) {
                    int bounce = this->sharcBounce;
                    ImGui::SliderInt("Min Query Bounce", &bounce, 0, 13);
                    ImGui::SetItemTooltip("Min bounce for SHARC query (update uses all bounces, aligned with RTXGI)");
                    if (bounce != this->sharcBounce) {
                        this->sharcBounce = bounce;
                    }
                    
                    int resScale = this->sharcUpdateResScale;
                    ImGui::SliderInt("Update Res Scale (1/N)", &resScale, 1, 10);
                    if (resScale != this->sharcUpdateResScale) {
                        this->sharcUpdateResScale = resScale;
                    }
                    
                    int maxBounces = this->sharcMaxBounces;
                    ImGui::SliderInt("Max Bounces", &maxBounces, 1, 32);
                    ImGui::SetItemTooltip("Max bounces for SHARC update pass.\nHigher = more accurate indirect lighting, but slower.");
                    if (maxBounces != this->sharcMaxBounces) {
                        this->sharcMaxBounces = maxBounces;
                    }
                    
                    int sampleThreshold = this->sharcSampleThreshold;
                    ImGui::SliderInt("Sample Threshold", &sampleThreshold, 1, 32);
                    if (sampleThreshold != this->sharcSampleThreshold) {
                        this->sharcSampleThreshold = sampleThreshold;
                    }
                    
                    ImGui::Separator();
                    bool compareMode = this->sharcCompareMode;
                    ImGui::Checkbox("Compare Mode (Cache vs GT)", &compareMode);
                    ImGui::SetItemTooltip("Show only bounces >= Min Query Bounce.\nSHARC ON: shows cached radiance\nSHARC OFF: shows ground truth for comparison");
                    if (compareMode != this->sharcCompareMode) {
                        this->sharcCompareMode = compareMode;
                        this->accumId = 0;
                    }
                    
                    bool showBeta = this->sharcShowBeta;
                    ImGui::Checkbox("Show Beta at Query Bounce", &showBeta);
                    ImGui::SetItemTooltip("Debug: Show throughput (beta) value at sharcBounce hit point.\nUseful for understanding energy loss before SHARC query.");
                    if (showBeta != this->sharcShowBeta) {
                        this->sharcShowBeta = showBeta;
                        this->accumId = 0;
                    }
                    
                    float queryScale = this->sharcQueryScale;
                    ImGui::SliderFloat("Query Scale", &queryScale, 0.0f, 2.0f, "%.2f");
                    ImGui::SetItemTooltip("Scale factor for SHARC query radiance.\n1.0 = normal, 0.0 = disable cache contribution (debug)");
                    if (queryScale != this->sharcQueryScale) {
                        this->sharcQueryScale = queryScale;
                        this->accumId = 0;
                    }
                    
                    float hitDistScale = this->sharcHitDistanceScale;
                    ImGui::SliderFloat("Hit Distance Scale", &hitDistScale, 0.0f, 3.0f, "%.2f");
                    ImGui::SetItemTooltip("Min hit distance = voxelSize * this value.\n1.732 = voxel diagonal (default), 0 = disable check");
                    if (hitDistScale != this->sharcHitDistanceScale) {
                        this->sharcHitDistanceScale = hitDistScale;
                        this->accumId = 0;
                    }
                    
                    float queryJitter = this->sharcQueryJitter;
                    ImGui::SliderFloat("Query Jitter", &queryJitter, 0.0f, 1.0f, "%.2f");
                    ImGui::SetItemTooltip("Jitter query position within voxel.\n0 = disabled (default), 1 = full voxel size");
                    if (queryJitter != this->sharcQueryJitter) {
                        this->sharcQueryJitter = queryJitter;
                        this->accumId = 0;
                    }

                    ImGui::Separator();
                    bool useHairTangent = this->sharcUseHairTangent;
                    ImGui::Checkbox("Use Hair Tangent", &useHairTangent);
                    ImGui::SetItemTooltip("For hair hits, use tangent-based SHARC key/query.\nOff = fall back to normal-based SHARC indexing.");
                    if (useHairTangent != this->sharcUseHairTangent) {
                        this->sharcUseHairTangent = useHairTangent;
                        this->resetSharc();
                        this->accumId = 0;
                    }

                    int hairQueryMode = this->sharcHairQueryMode;
                    ImGui::SliderInt("Hair Query Mode", &hairQueryMode, 0, 2);
                    ImGui::SetItemTooltip("0 = Off\n1 = Hard validity gate\n2 = Blend nearby direction-aware cache entries");
                    if (hairQueryMode != this->sharcHairQueryMode) {
                        this->sharcHairQueryMode = hairQueryMode;
                        this->resetSharc();
                        this->accumId = 0;
                    }

                    if (this->sharcHairQueryMode > 0) {
                        float tangentThreshold = this->sharcHairTangentThreshold;
                        ImGui::SliderFloat("Tangent Threshold", &tangentThreshold, 0.0f, 1.0f, "%.2f");
                        ImGui::SetItemTooltip("Require abs(dot(t_query, t_entry)) >= threshold.");
                        if (tangentThreshold != this->sharcHairTangentThreshold) {
                            this->sharcHairTangentThreshold = tangentThreshold;
                            this->resetSharc();
                            this->accumId = 0;
                        }

                        float parallelScale = this->sharcHairParallelScale;
                        ImGui::SliderFloat("Parallel Scale", &parallelScale, 0.25f, 6.0f, "%.2f");
                        ImGui::SetItemTooltip("Support radius along hair tangent, in voxel units.");
                        if (parallelScale != this->sharcHairParallelScale) {
                            this->sharcHairParallelScale = parallelScale;
                            this->resetSharc();
                            this->accumId = 0;
                        }

                        float perpScale = this->sharcHairPerpScale;
                        ImGui::SliderFloat("Perp Scale", &perpScale, 0.10f, 3.0f, "%.2f");
                        ImGui::SetItemTooltip("Support radius across hair tangent, in voxel units.");
                        if (perpScale != this->sharcHairPerpScale) {
                            this->sharcHairPerpScale = perpScale;
                            this->resetSharc();
                            this->accumId = 0;
                        }
                    }
                    
                    ImGui::TreePop();
                }
                
                // Resolve settings
                if (ImGui::TreeNode("Resolve Settings")) {
                    int staleMax = this->sharcStaleFrameMax;
                    ImGui::SliderInt("Stale Frame Max", &staleMax, 1, 64);
                    if (staleMax != this->sharcStaleFrameMax) {
                        this->sharcStaleFrameMax = staleMax;
                    }
                    
                    int accumFrames = this->sharcAccumFrameNum;
                    ImGui::SliderInt("Accumulation Frames", &accumFrames, 1, 128);
                    if (accumFrames != this->sharcAccumFrameNum) {
                        this->sharcAccumFrameNum = accumFrames;
                    }

                    bool blendAdjacent = this->sharcBlendAdjacentLevels;
                    ImGui::Checkbox("Blend Adjacent Levels", &blendAdjacent);
                    if (blendAdjacent != this->sharcBlendAdjacentLevels) {
                        this->sharcBlendAdjacentLevels = blendAdjacent;
                    }
                    
                    float radScale = this->sharcRadianceScale;
                    ImGui::SliderFloat("Radiance Scale", &radScale, 0.1f, 100.f, "%.1f");
                    if (radScale != this->sharcRadianceScale) {
                        this->sharcRadianceScale = radScale;
                    }
                    
                    bool antiFirefly = this->sharcAntiFirefly;
                    ImGui::Checkbox("Anti-Firefly Filter", &antiFirefly);
                    if (antiFirefly != this->sharcAntiFirefly) {
                        this->sharcAntiFirefly = antiFirefly;
                    }
                    
                    bool disableEnvMiss = this->sharcDisableEnvMiss;
                    ImGui::Checkbox("Disable Env Miss Update", &disableEnvMiss);
                    ImGui::SetItemTooltip("Skip SharcUpdateMiss when ray escapes to environment.\nPrevents environment radiance from diluting hair color in cache.");
                    if (disableEnvMiss != this->sharcDisableEnvMiss) {
                        this->sharcDisableEnvMiss = disableEnvMiss;
                    }
                    
                    bool clearEveryFrame = this->sharcClearEveryFrame;
                    ImGui::Checkbox("Clear Every Frame", &clearEveryFrame);
                    ImGui::SetItemTooltip("Reset cache every frame (no temporal accumulation).\nUseful for debugging single-frame cache behavior.");
                    if (clearEveryFrame != this->sharcClearEveryFrame) {
                        this->sharcClearEveryFrame = clearEveryFrame;
                    }
                    
                    bool materialDemod = this->sharcMaterialDemod;
                    ImGui::Checkbox("Material Demodulation", &materialDemod);
                    ImGui::SetItemTooltip("Demodulate material color to preserve color details.\nFor hair: uses exp(-sig_a * scale) as demodulation factor.");
                    if (materialDemod != this->sharcMaterialDemod) {
                        this->sharcMaterialDemod = materialDemod;
                        this->accumId = 0;
                    }
                    
                    if (this->sharcMaterialDemod) {
                        float demodScale = this->sharcDemodScale;
                        ImGui::SliderFloat("Demod Scale", &demodScale, 0.5f, 5.0f, "%.1f");
                        ImGui::SetItemTooltip("Scale for hair demodulation: exp(-sig_a * scale).\nHigher = stronger color preservation.");
                        if (demodScale != this->sharcDemodScale) {
                            this->sharcDemodScale = demodScale;
                            this->accumId = 0;
                        }
                    }
                    
                    ImGui::TreePop();
                }
                
                // Grid settings
                if (ImGui::TreeNode("Grid Settings")) {
                    float sceneScale = this->sharcSceneScale;
                    ImGui::SliderFloat("Scene Scale (0=auto)", &sceneScale, 0.f, 1000.f);
                    if (sceneScale != this->sharcSceneScale) {
                        this->sharcSceneScale = sceneScale;
                    }
                    
                    float logBase = this->sharcLogBase;
                    ImGui::SliderFloat("Log Base", &logBase, 1.5f, 4.f);
                    if (logBase != this->sharcLogBase) {
                        this->sharcLogBase = logBase;
                    }
                    
                    int levelBias = this->sharcLevelBias;
                    ImGui::SliderInt("Level Bias", &levelBias, -5, 5);
                    if (levelBias != this->sharcLevelBias) {
                        this->sharcLevelBias = levelBias;
                    }
                    
                    ImGui::TreePop();
                }
                
                ImGui::Separator();
                
                if (ImGui::Button("Reset Cache")) {
                    resetSharc();
                    this->cameraChanged();
                }
                
                if (sharcInitialized) {
                    int updateW = max(1, this->fbSize.x / max(1, this->sharcUpdateResScale));
                    int updateH = max(1, this->fbSize.y / max(1, this->sharcUpdateResScale));
                    ImGui::Text("Update: %dx%d, Full: %dx%d", updateW, updateH, this->fbSize.x, this->fbSize.y);
                    ImGui::Text("Cache: %u entries (%.1f MB)", sharcCapacity,
                               float(sharcCapacity * (sizeof(uint64_t) + sizeof(uint32_t) + 
                               sizeof(SharcAccumulationData) + sizeof(SharcPackedData) +
                               sizeof(SharcHairMetaData))) / (1024.0f * 1024.0f));
                }
            }
        }

        // GPU Profiler
        if (ImGui::CollapsingHeader("Profiler")) {
            ImGui::Checkbox("Enable GPU Profiler", &this->profilerEnabled);
            ImGui::SetItemTooltip("Measures GPU time for each rendering pass using CUDA events.\nNote: synchronization may slightly reduce throughput.");
            
            if (profilerEnabled && profilerInitialized) {
                ImGui::Separator();
                
                const float barWidth = 200.f;
                
                // Total frame GPU time
                ImGui::Text("Total GPU:       %6.2f ms", profTotalMs);
                
                // Per-pass breakdown
                if (sharcEnabled && sharcUpdateEnabled) {
                    ImGui::Text("  Update Pass:   %6.2f ms", profUpdateMs);
                    float updateFrac = (profTotalMs > 0.f) ? profUpdateMs / profTotalMs : 0.f;
                    ImGui::ProgressBar(updateFrac, ImVec2(barWidth, 0), "");
                    ImGui::SameLine();
                    ImGui::Text("%.0f%%", updateFrac * 100.f);
                    
                    ImGui::Text("  Resolve Pass:  %6.2f ms", profResolveMs);
                    float resolveFrac = (profTotalMs > 0.f) ? profResolveMs / profTotalMs : 0.f;
                    ImGui::ProgressBar(resolveFrac, ImVec2(barWidth, 0), "");
                    ImGui::SameLine();
                    ImGui::Text("%.0f%%", resolveFrac * 100.f);
                }
                
                ImGui::Text("  Render Pass:   %6.2f ms", profRenderMs);
                float renderFrac = (profTotalMs > 0.f) ? profRenderMs / profTotalMs : 0.f;
                ImGui::ProgressBar(renderFrac, ImVec2(barWidth, 0), "");
                ImGui::SameLine();
                ImGui::Text("%.0f%%", renderFrac * 100.f);
                
                ImGui::Separator();
                ImGui::Text("GPU FPS: %.1f", (profTotalMs > 0.f) ? 1000.f / profTotalMs : 0.f);
            }
        }

        // All lighting controls
        if (ImGui::CollapsingHeader("Light")) {

            if (ImGui::TreeNode("Environment Light")) {
                // Controls for environment light
                if (this->currentScene.hasEnvLight) {
                    float scale = this->currentScene.envScale;
                    ImGui::SliderFloat("Radiance Scale", &scale, 0.f, 4.f);
                    if (scale != this->currentScene.envScale) {
                        this->currentScene.envScale = scale;
                        owlParamsSet1f(this->launchParams, "envScale", scale);
                        this->cameraChanged();
                    }

                    float rot = this->currentScene.envRotPhi;
                    ImGui::SliderFloat("Rotation", &rot, 0.f, 3.14159f * 2.f);
                    if (rot != this->currentScene.envRotPhi) {
                        this->currentScene.envRotPhi = rot;
                        owlParamsSet1f(this->launchParams, "envRotPhi", rot);
                        this->cameraChanged();
                    }
                }

                ImGui::TreePop();
            }

            if (ImGui::TreeNode("Directional Lights")) {
                // Controls for directional lights
                for (int i = 0; i < this->num_dlights; i++) {
                    float dLightFrom[3] = {
                        this->dLightList[i].from.x,
                        this->dLightList[i].from.y,
                        this->dLightList[i].from.z
                    };
                    std::string sname = "Dir " + std::to_string(i + 1);
                    ImGui::SliderFloat3(sname.c_str(), dLightFrom, -1.f, 1.f);

                    if (dLightFrom[0] != this->dLightList[i].from.x ||
                        dLightFrom[1] != this->dLightList[i].from.y ||
                        dLightFrom[2] != this->dLightList[i].from.z) {

                        this->dLightList[i].from = vec3f(dLightFrom[0], dLightFrom[1], dLightFrom[2]);
                        this->uploadLights();
                    }

                    float emit[3] = {
                        this->dLightList[i].emit.x,
                        this->dLightList[i].emit.y,
                        this->dLightList[i].emit.z
                    };
                    std::string ename = "Emit " + std::to_string(i + 1);
                    ImGui::SliderFloat3(ename.c_str(), emit, 0.f, 10.f);
                    if (emit[0] != this->dLightList[i].emit.x ||
                        emit[1] != this->dLightList[i].emit.y ||
                        emit[2] != this->dLightList[i].emit.z) {
                        this->dLightList[i].emit = vec3f(emit[0], emit[1], emit[2]);
                        this->uploadLights();
                    }
                }

                ImGui::TreePop();
            }
        }

        // Controls specific to path integrators
        if (ImGui::CollapsingHeader("Path tracing")) {
            int v2 = this->currentScene.path_v2;
            ImGui::SliderInt("Path Depth v2", &v2, 1, 50);

            int v1 = this->currentScene.path_v1;
            ImGui::SliderInt("Path Depth v1", &v1, 1, v2);

            if (v1 != this->currentScene.path_v1 || v2 != this->currentScene.path_v2) {
                this->currentScene.path_v1 = v1;
                this->currentScene.path_v2 = v2;

                owlParamsSet1i(this->launchParams, "pathV1", v1);
                owlParamsSet1i(this->launchParams, "pathV2", v2);

                this->cameraChanged();
            }
        }

        if (ImGui::CollapsingHeader("Camera")) {
            float camSpeed = this->camera.motionSpeed;
            if (ImGui::SliderFloat("Camera Speed", &camSpeed, 0.01f, 500.0f, "%.3f", ImGuiSliderFlags_Logarithmic)) {
                this->camera.motionSpeed = camSpeed;
            }
            if (ImGui::Button("Copy Camera JSON")) {
                copyCameraJsonToClipboard();
            }
        }

        if (ImGui::Button("Save PNG")) {
            char indexedName[256];
            std::string basePath = this->currentScene.renderOutput;
            size_t dotPos = basePath.rfind('.');
            std::string stem = (dotPos != std::string::npos) ? basePath.substr(0, dotPos) : basePath;
            std::string ext = (dotPos != std::string::npos) ? basePath.substr(dotPos) : ".png";
            snprintf(indexedName, sizeof(indexedName), "%s%03d%s", stem.c_str(), this->saveIndex, ext.c_str());
            this->screenShot(std::string(indexedName));
            this->saveIndex++;
        }

        if (ImGui::Button("Save EXR")) {
            char indexedName[256];
            std::string basePath = this->currentScene.renderOutput;
            size_t dotPos = basePath.rfind('.');
            std::string stem = (dotPos != std::string::npos) ? basePath.substr(0, dotPos) : basePath;
            snprintf(indexedName, sizeof(indexedName), "%s%03d.exr", stem.c_str(), this->saveIndex);
            this->screenShotEXR(std::string(indexedName));
            this->saveIndex++;
        }

        ImGui::End();
    }

    ImGui::Render();
    ImGui_ImplOpenGL2_RenderDrawData(ImGui::GetDrawData());
}

// ====================================================
// Mostly constant functions
// ====================================================
void RenderWindowPT::screenShotEXR(std::string fname)
{
    int bufferSize = this->fbSize.x * this->fbSize.y * 4 * sizeof(float);
    float* hostBuffer = (float*)malloc(bufferSize);
    cudaMemcpy(hostBuffer, this->averageBuffer, bufferSize, cudaMemcpyDeviceToHost);
    saveEXR(fname, hostBuffer, this->accumId, this->fbSize.x, this->fbSize.y);
    free(hostBuffer);
}

RenderWindowPT::RenderWindowPT(Scene& scene, vec2i resolution, bool interactive)
    : owl::viewer::OWLViewer("Real-time hair", resolution, interactive, false)
{
    this->currentScene = scene;

    this->camera.setOrientation(scene.camera.from,
        scene.camera.at,
        scene.camera.up,
        owl::viewer::toDegrees(acosf(scene.camera.cosFovy)));

    if (interactive)
        this->enableFlyMode();

    // Initialize IMGUI
    IMGUI_CHECKVERSION();
    ImGui::CreateContext();
    ImGuiIO& io = ImGui::GetIO(); (void)io;

    ImGui::StyleColorsDark();

    ImGui_ImplGlfw_InitForOpenGL(this->handle, true);
    ImGui_ImplOpenGL2_Init();

    // Context & Module creation, accumulation buffer initialize
    context = owlContextCreate(nullptr, 1);
    module = owlModuleCreate(context, path_tracing_ptx);

    owlContextSetRayTypeCount(context, NUM_RAY_TYPES);
    owlEnableCurves(context); // Required since hair strands are represented by curves
}

std::string RenderWindowPT::cameraJsonSnippet() const
{
    SceneCamera cam;
    cam.from = this->camera.getFrom();
    cam.at = this->camera.getAt();
    cam.up = this->camera.getUp();
    cam.cosFovy = this->camera.getCosFovy();

    std::ostringstream ss;
    ss << std::fixed << std::setprecision(9);
    ss << "\"camera\": {\n";
    ss << "  \"cos_fovy\": " << cam.cosFovy << ",\n";
    ss << "  \"from\": [\n";
    ss << "    " << cam.from.x << ",\n";
    ss << "    " << cam.from.y << ",\n";
    ss << "    " << cam.from.z << "\n";
    ss << "  ],\n";
    ss << "  \"to\": [\n";
    ss << "    " << cam.at.x << ",\n";
    ss << "    " << cam.at.y << ",\n";
    ss << "    " << cam.at.z << "\n";
    ss << "  ],\n";
    ss << "  \"up\": [\n";
    ss << "    " << cam.up.x << ",\n";
    ss << "    " << cam.up.y << ",\n";
    ss << "    " << cam.up.z << "\n";
    ss << "  ]\n";
    ss << "},";
    return ss.str();
}

void RenderWindowPT::copyCameraJsonToClipboard()
{
    glfwSetClipboardString(this->handle, cameraJsonSnippet().c_str());
    LOG("[Camera] Copied camera JSON to clipboard (UI button).");
}

void RenderWindowPT::customKey(char key, const vec2i& pos)
{
    if (key == '1' || key == '!') {
        this->camera.setOrientation(this->camera.getFrom(), vec3f(0.f), vec3f(0.f, 0.f, 1.f), this->camera.getFovyInDegrees());
        this->cameraChanged();
    }
    else if (key == 'R' || key == 'r') {
        SceneCamera cam;
        cam.from = this->camera.getFrom();
        cam.at = this->camera.getAt();
        cam.up = this->camera.getUp();
        cam.cosFovy = this->camera.getCosFovy();

        nlohmann::json oneCameraJson;
        std::vector<float> from, at, up;

        for (int i = 0; i < 3; i++) {
            from.push_back(cam.from[i]);
            at.push_back(cam.at[i]);
            up.push_back(cam.up[i]);
        }

        oneCameraJson["from"] = from;
        oneCameraJson["to"] = at;
        oneCameraJson["up"] = up;
        oneCameraJson["cos_fovy"] = cam.cosFovy;

        this->currentScene.json["camera"] = oneCameraJson;
        std::ofstream outputFile(this->currentScene.jsonFilePath);
        outputFile << std::setw(4) << this->currentScene.json << std::endl;
    }
    else if (key == 'P' || key == 'p') {
        SceneCamera cam;
        cam.from = this->camera.getFrom();
        cam.at = this->camera.getAt();
        cam.up = this->camera.getUp();
        cam.cosFovy = this->camera.getCosFovy();

        this->currentScene.cameraPath.push_back(cam);
    }
    else if (key == '[' || key == '{') {
        this->currentScene.cameraPath.clear();
    }
    else if (key == 'o' || key == 'O') {
        std::vector<nlohmann::json> cameraPathJson;
        
        for (auto cam : this->currentScene.cameraPath) {
            nlohmann::json oneCameraJson;
            std::vector<float> from, at, up;

            for (int i = 0; i < 3; i++) {
                from.push_back(cam.from[i]);
                at.push_back(cam.at[i]);
                up.push_back(cam.up[i]);
            }

            oneCameraJson["from"] = from;
            oneCameraJson["to"] = at;
            oneCameraJson["up"] = up;
            oneCameraJson["cos_fovy"] = cam.cosFovy;

            cameraPathJson.push_back(oneCameraJson);
        }

        this->currentScene.json["camera_path"] = cameraPathJson;
        std::ofstream outputFile(this->currentScene.jsonFilePath);
        outputFile << std::setw(4) << this->currentScene.json << std::endl;
    }
}

void RenderWindowPT::cameraChanged()
{
    // Reset accumulation buffer, to restart MC sampling
    this->accumId = 0;

    const vec3f lookFrom = camera.getFrom();
    const vec3f lookAt = camera.getAt();
    const vec3f lookUp = camera.getUp();
    const float cosFovy = camera.getCosFovy();

    // ----------- compute variable values  ------------------
    vec3f camera_pos = lookFrom;
    vec3f camera_d00
        = normalize(lookAt - lookFrom);
    float aspect = fbSize.x / float(fbSize.y);
    vec3f camera_ddu
        = cosFovy * aspect * normalize(cross(camera_d00, lookUp));
    vec3f camera_ddv
        = cosFovy * normalize(cross(camera_ddu, camera_d00));
    camera_d00 -= 0.5f * camera_ddu;
    camera_d00 -= 0.5f * camera_ddv;

    // ----------- set variables  ----------------------------
    owlRayGenSet1ul(this->rayGen, "frameBuffer", (uint64_t)this->fbPointer);
    owlRayGenSet2i(this->rayGen, "frameBufferSize", (const owl2i&)this->fbSize);

    owlParamsSet3f(this->launchParams, "camera.pos", (const owl3f&)camera_pos);
    owlParamsSet3f(this->launchParams, "camera.dir_00", (const owl3f&)camera_d00);
    owlParamsSet3f(this->launchParams, "camera.dir_du", (const owl3f&)camera_ddu);
    owlParamsSet3f(this->launchParams, "camera.dir_dv", (const owl3f&)camera_ddv);

    this->sbtDirty = true;
}

void RenderWindowPT::uploadLights()
{
    // Directional lights
    this->dLightsBuffer = owlDeviceBufferCreate(context,
        OWL_USER_TYPE(DirectionalLight), this->dLightList.size(), this->dLightList.data());

    // Upload directional lights
    owlParamsSetBuffer(this->launchParams, "dLights", dLightsBuffer);
    owlParamsSet1i(this->launchParams, "num_dlights", num_dlights);

    this->cameraChanged();
}

// ---- PT overrides --------------------------------------------------------
struct PtTaskParams {
    int   MIS         = -1;   // -1 = not set
    int   env_pdf     = -1;
    int   path_v1     = -1;
    int   path_v2     = -1;
    float env_scale   = -1.f;
    float env_rot_phi = -1.f;
};

// ---- Hair overrides ------------------------------------------------------
struct HairTaskParams {
    bool  has_sigma_a = false;
    vec3f sigma_a = vec3f(0.f);
};

// ---- Environment overrides ----------------------------------------------
struct EnvTaskParams {
    std::string exr = "";
    float scale = -1.f;
    float rotation = -1.f;
};

// ---- SHARC overrides -----------------------------------------------------
struct SharcTaskParams {
    int   debug_mode         = -1;
    int   enabled            = -1;
    int   update_enabled     = -1;
    int   capacity           = -1;
    int   bounce             = -1;
    int   max_bounces        = -1;
    int   update_res_scale   = -1;
    int   sample_threshold   = -1;
    float query_scale        = -1.f;
    float hit_distance_scale = -1.f;
    float query_jitter       = -1.f;
    int   stale_frame_max    = -1;
    int   accum_frame_num    = -1;
    float radiance_scale     = -1.f;
    int   anti_firefly       = -1;
    int   disable_env_miss   = -1;
    int   clear_every_frame  = -1;
    int   material_demod     = -1;
    float demod_scale        = -1.f;
    float scene_scale        = -1.f;
    float log_base           = -1.f;
    int   level_bias         = -1;
    int   use_hair_tangent   = -1;
    int   hair_query_mode    = -1;
    float hair_tangent_threshold = -1.f;
    float hair_parallel_scale = -1.f;
    float hair_perp_scale     = -1.f;
    int   blend_adjacent_levels = -1;
};

// ---- Top-level task ------------------------------------------------------
struct RenderTask {
    bool            valid      = false;
    std::string     scene      = "";
    int             spp        = -1;
    std::string     output     = "";
    bool            batch_mode = true; // task mode default: auto-exit when finished
    std::string     task_file  = "";
    nlohmann::json  raw_json   = nlohmann::json::object();
    std::vector<int> snapshot_times_ms;
    std::vector<int> snapshot_spps;
    PtTaskParams    pt;
    HairTaskParams  hair;
    EnvTaskParams   env;
    SharcTaskParams sharc;
};

RenderTask parseRenderTask(const std::string& taskFile)
{
    RenderTask task;
    task.task_file = taskFile;
    std::ifstream f(taskFile);
    if (!f.is_open()) {
        std::cerr << "[Task] Cannot open: " << taskFile << "\n";
        return task;
    }
    nlohmann::json j;
    try { f >> j; }
    catch (const std::exception& e) {
        std::cerr << "[Task] JSON error: " << e.what() << "\n";
        return task;
    }
    task.raw_json = j;

    // basic
    if (j.contains("scene"))      task.scene      = j["scene"].get<std::string>();
    if (j.contains("spp"))        task.spp        = j["spp"].get<int>();
    if (j.contains("output"))     task.output     = j["output"].get<std::string>();
    if (j.contains("batch_mode")) task.batch_mode = j["batch_mode"].get<bool>();
    if (j.contains("snapshot_times_ms") && j["snapshot_times_ms"].is_array()) {
        for (auto& v : j["snapshot_times_ms"]) {
            try {
                int t = v.get<int>();
                if (t > 0) task.snapshot_times_ms.push_back(t);
            } catch (...) {}
        }
        std::sort(task.snapshot_times_ms.begin(), task.snapshot_times_ms.end());
        task.snapshot_times_ms.erase(
            std::unique(task.snapshot_times_ms.begin(), task.snapshot_times_ms.end()),
            task.snapshot_times_ms.end());
    }
    if (j.contains("snapshot_spps") && j["snapshot_spps"].is_array()) {
        for (auto& v : j["snapshot_spps"]) {
            try {
                int s = v.get<int>();
                if (s > 0) task.snapshot_spps.push_back(s);
            } catch (...) {}
        }
        std::sort(task.snapshot_spps.begin(), task.snapshot_spps.end());
        task.snapshot_spps.erase(
            std::unique(task.snapshot_spps.begin(), task.snapshot_spps.end()),
            task.snapshot_spps.end());
    }

    // PT params
    if (j.contains("pt")) {
        auto& p = j["pt"];
        if (p.contains("MIS"))         task.pt.MIS         = p["MIS"].get<bool>() ? 1 : 0;
        if (p.contains("env_pdf"))     task.pt.env_pdf     = p["env_pdf"].get<bool>() ? 1 : 0;
        if (p.contains("path_v1"))     task.pt.path_v1     = p["path_v1"].get<int>();
        if (p.contains("path_v2"))     task.pt.path_v2     = p["path_v2"].get<int>();
        if (p.contains("env_scale"))   task.pt.env_scale   = p["env_scale"].get<float>();
        if (p.contains("env_rot_phi")) task.pt.env_rot_phi = p["env_rot_phi"].get<float>();
    }

    // Hair params
    if (j.contains("hair")) {
        auto& h = j["hair"];
        if (h.contains("sigma_a") && h["sigma_a"].is_array() && h["sigma_a"].size() >= 3) {
            task.hair.sigma_a = vec3f(
                h["sigma_a"][0].get<float>(),
                h["sigma_a"][1].get<float>(),
                h["sigma_a"][2].get<float>());
            task.hair.has_sigma_a = true;
        }
    }

    // Environment params
    if (j.contains("env")) {
        auto& e = j["env"];
        if (e.contains("exr"))      task.env.exr = e["exr"].get<std::string>();
        if (e.contains("scale"))    task.env.scale = e["scale"].get<float>();
        if (e.contains("rotation")) task.env.rotation = e["rotation"].get<float>();
    }

    // SHARC params
    if (j.contains("sharc")) {
        auto& s = j["sharc"];
#define READ_BOOL(field)  if (s.contains(#field)) task.sharc.field = s[#field].get<bool>() ? 1 : 0
#define READ_INT(field)   if (s.contains(#field)) task.sharc.field = s[#field].get<int>()
#define READ_FLOAT(field) if (s.contains(#field)) task.sharc.field = s[#field].get<float>()
        READ_BOOL(enabled);
        READ_INT(debug_mode);
        READ_BOOL(update_enabled);
        READ_INT(capacity);
        READ_INT(bounce);
        READ_INT(max_bounces);
        READ_INT(update_res_scale);
        READ_INT(sample_threshold);
        READ_FLOAT(query_scale);
        READ_FLOAT(hit_distance_scale);
        READ_FLOAT(query_jitter);
        READ_INT(stale_frame_max);
        READ_INT(accum_frame_num);
        READ_FLOAT(radiance_scale);
        READ_BOOL(anti_firefly);
        READ_BOOL(disable_env_miss);
        READ_BOOL(clear_every_frame);
        READ_BOOL(material_demod);
        READ_FLOAT(demod_scale);
        READ_FLOAT(scene_scale);
        READ_FLOAT(log_base);
        READ_INT(level_bias);
        READ_BOOL(use_hair_tangent);
        READ_INT(hair_query_mode);
        READ_FLOAT(hair_tangent_threshold);
        READ_FLOAT(hair_parallel_scale);
        READ_FLOAT(hair_perp_scale);
        READ_BOOL(blend_adjacent_levels);
#undef READ_BOOL
#undef READ_INT
#undef READ_FLOAT
    }

    task.valid = true;
    std::cout << "[Task] Loaded: scene=" << task.scene
              << "  spp=" << task.spp
              << "  output=" << task.output
              << "  sharc=" << (task.sharc.enabled == 1 ? "ON" : task.sharc.enabled == 0 ? "OFF" : "default")
              << "\n";
    return task;
}

static void applyTaskToScene(Scene& scene, const RenderTask& task)
{
    if (!task.valid) return;

    if (task.spp > 0)         scene.spp = task.spp;
    if (!task.output.empty()) scene.renderOutput = task.output;

    // Backward-compatible env overrides under "pt".
    if (task.pt.env_scale >= 0.f)   scene.envScale = task.pt.env_scale;
    if (task.pt.env_rot_phi >= 0.f) scene.envRotPhi = task.pt.env_rot_phi;

    // Hair color override
    if (task.hair.has_sigma_a) {
        scene.sig_a = task.hair.sigma_a;
        LOG("[Task] Override hair.sigma_a");
    }

    // Environment override (preferred location: top-level "env")
    if (task.env.scale >= 0.f)    scene.envScale = task.env.scale;
    if (task.env.rotation >= 0.f) scene.envRotPhi = task.env.rotation;
    if (!task.env.exr.empty()) {
        std::string path = task.env.exr;
        if (!loadEnvTexture(path, &scene.env)) {
            LOG("[Task] Failed to load env.exr from task: " + path);
        } else {
            scene.hasEnvLight = true;
            LOG("[Task] Override env.exr -> " + path);
        }
    }
}

static std::string ensureExrPath(std::string out)
{
    if (out.size() >= 4 && out.substr(out.size() - 4) == ".exr")
        return out;
    size_t dot = out.rfind('.');
    return (dot != std::string::npos) ? out.substr(0, dot) + ".exr" : out + ".exr";
}

static std::string taskMetaPathFromExr(const std::string& exrPath)
{
    size_t dot = exrPath.rfind('.');
    if (dot == std::string::npos) return exrPath + ".meta.json";
    return exrPath.substr(0, dot) + ".meta.json";
}

static std::string appendSuffixBeforeExt(const std::string& path, const std::string& suffix)
{
    size_t dot = path.rfind('.');
    if (dot == std::string::npos) return path + suffix;
    return path.substr(0, dot) + suffix + path.substr(dot);
}

static void writeTaskMetaJson(
    const std::string& exrPath,
    const RenderTask& task,
    const RenderWindowPT& win,
    int warmupFrames,
    int countStartAccum)
{
    nlohmann::json meta;
    meta["exr_output"] = exrPath;
    meta["gpu_render_time_ms"] = win.taskGpuRenderMsTotal;
    meta["gpu_render_time_last_frame_ms"] = win.taskGpuRenderMsLast;
    meta["spp_rendered"] = win.accumId - countStartAccum;
    meta["total_frames_rendered"] = win.accumId;
    meta["warmup_frames"] = warmupFrames;
    meta["task_json_path"] = task.task_file;
    meta["task_json"] = task.raw_json;

    const std::string metaPath = taskMetaPathFromExr(exrPath);
    std::ofstream out(metaPath);
    if (!out.is_open()) {
        LOG("[Task] Failed to write meta json: " + metaPath);
        return;
    }
    out << std::setw(2) << meta << std::endl;
    LOG("[Task] Wrote meta json -> " + metaPath);
}

void applyTaskToWindow(RenderWindowPT& win, const RenderTask& task)
{
    if (!task.valid) return;

    // PT params: 只写成员变量，initialize() 里统一 owlParamsSet
    auto& p = task.pt;
    if (p.MIS         >= 0) win.currentScene.MIS            = p.MIS;
    if (p.env_pdf     >= 0) win.currentScene.envPdfSampling = p.env_pdf;
    if (p.path_v1     >= 0) win.currentScene.path_v1        = p.path_v1;
    if (p.path_v2     >= 0) win.currentScene.path_v2        = p.path_v2;
    if (p.env_scale   >= 0) win.currentScene.envScale       = p.env_scale;
    if (p.env_rot_phi >= 0) win.currentScene.envRotPhi      = p.env_rot_phi;

    // SHARC params: 只写成员变量，render loop 里统一上传
    auto& s = task.sharc;
    if (s.debug_mode        >= 0) win.sharcDebugMode        = s.debug_mode;
    if (s.capacity          > 0)  win.sharcCapacity        = (uint32_t)s.capacity;
    if (s.bounce            >= 0) win.sharcBounce           = s.bounce;
    if (s.max_bounces       >= 0) win.sharcMaxBounces       = s.max_bounces;
    if (s.update_res_scale  >= 0) win.sharcUpdateResScale   = s.update_res_scale;
    if (s.sample_threshold  >= 0) win.sharcSampleThreshold  = s.sample_threshold;
    if (s.query_scale       >= 0) win.sharcQueryScale       = s.query_scale;
    if (s.hit_distance_scale>= 0) win.sharcHitDistanceScale = s.hit_distance_scale;
    if (s.query_jitter      >= 0) win.sharcQueryJitter      = s.query_jitter;
    if (s.stale_frame_max   >= 0) win.sharcStaleFrameMax    = s.stale_frame_max;
    if (s.accum_frame_num   >= 0) win.sharcAccumFrameNum    = s.accum_frame_num;
    if (s.radiance_scale    >= 0) win.sharcRadianceScale    = s.radiance_scale;
    if (s.anti_firefly      >= 0) win.sharcAntiFirefly      = s.anti_firefly;
    if (s.disable_env_miss  >= 0) win.sharcDisableEnvMiss   = s.disable_env_miss;
    if (s.clear_every_frame >= 0) win.sharcClearEveryFrame  = s.clear_every_frame;
    if (s.material_demod    >= 0) win.sharcMaterialDemod    = s.material_demod;
    if (s.demod_scale       >= 0) win.sharcDemodScale       = s.demod_scale;
    if (s.scene_scale       >= 0) win.sharcSceneScale       = s.scene_scale;
    if (s.log_base          >= 0) win.sharcLogBase          = s.log_base;
    if (s.level_bias        >= 0) win.sharcLevelBias        = s.level_bias;
    if (s.use_hair_tangent  >= 0) win.sharcUseHairTangent   = s.use_hair_tangent;
    if (s.hair_query_mode >= 0) win.sharcHairQueryMode = s.hair_query_mode;
    if (s.hair_tangent_threshold >= 0) win.sharcHairTangentThreshold = s.hair_tangent_threshold;
    if (s.hair_parallel_scale >= 0) win.sharcHairParallelScale = s.hair_parallel_scale;
    if (s.hair_perp_scale >= 0) win.sharcHairPerpScale = s.hair_perp_scale;
    if (s.blend_adjacent_levels >= 0) win.sharcBlendAdjacentLevels = s.blend_adjacent_levels;
    if (s.update_enabled    >= 0) win.sharcUpdateEnabled    = s.update_enabled;
    if (s.enabled           >= 0) win.sharcEnabled          = s.enabled;
}

// ====================================================
// Entry point
// ====================================================
int main(int argc, char** argv)
{
    std::string renderer = "Path Tracing";

    // ------------------------------------------------------------------
    // Argument parsing:
    //   argv[1]          -> scene config path  (original behaviour)
    //   --task <file>    -> load scene / spp / output from task JSON
    // ------------------------------------------------------------------
    std::string currentScene = "F:/code/HairMSNN_code2/scenes/curly/config.json"; 
    RenderTask task;

    for (int i = 1; i < argc; i++) {
        std::string arg = argv[i];
        if (arg == "--task" && i + 1 < argc) {
            task = parseRenderTask(argv[++i]);
        } else {
            currentScene = arg;
        }
    }

    if (task.valid && !task.scene.empty())
        currentScene = task.scene;

    LOG("Loading scene " + currentScene);

    Scene scene;
    bool success = parseScene(currentScene, scene);
    if (!success) {
        LOG("Error loading scene");
        return -1;
    }

    // Task overrides scene settings (hair/env/output/spp)
    applyTaskToScene(scene, task);

    vec2i resolution(scene.imgWidth, scene.imgHeight);
    RenderWindowPT win(scene, resolution, true);
    win.setTitle(renderer);

    // ------------------------------------------------------------------
    // Apply task overrides to window (SHARC params, PT params, etc.)
    // Must be called after win is fully constructed & initialized.
    // ------------------------------------------------------------------
    if (task.valid)
        applyTaskToWindow(win, task);
    if (task.valid)
        win.taskTimingEnabled = true;

    // ------------------------------------------------------------------
    // Task mode: auto-save EXR once accumId reaches spp.
    //   batch_mode=false  -> window stays open (interactive)
    //   batch_mode=true   -> window closes after saving (Python batch)
    // ------------------------------------------------------------------
    if (task.valid) {
        bool saved = false;
        bool snapshotMode = !task.snapshot_times_ms.empty();
        size_t nextSnapshotIdx = 0;
        bool snapshotSppMode = !task.snapshot_spps.empty();
        size_t nextSnapshotSppIdx = 0;
        const bool sharcTask = (task.sharc.enabled == 1) || win.sharcEnabled;
        const int warmupFrames = sharcTask ? 300 : 0;
        int countStartAccum = 0;
        bool warmupDone = (warmupFrames == 0);

        if (warmupFrames > 0) {
            LOG("[Task] SHARC warmup enabled: " + std::to_string(warmupFrames) + " frames");
        }

        
        win.showAndRun([&]() -> bool {
            if (!warmupDone && win.accumId >= warmupFrames) {
                warmupDone = true;
                // Warmup ends: keep SHARC cache, but restart image accumulation
                // so final EXR contains only post-warmup spp.
                win.accumId = 0;
                countStartAccum = 0;
                // Reset GPU-only task timer after SHARC cache warmup.
                win.taskGpuRenderMsTotal = 0.0;
                win.taskGpuRenderMsLast = 0.f;
                LOG("[Task] Warmup finished. Reset image accumulation; start counting post-warmup spp.");
            }

            const int countedSpp = win.accumId - countStartAccum;
            if (warmupDone && snapshotMode) {
                while (nextSnapshotIdx < task.snapshot_times_ms.size() &&
                       win.taskGpuRenderMsTotal >= double(task.snapshot_times_ms[nextSnapshotIdx])) {
                    std::string outBase = ensureExrPath(win.currentScene.renderOutput);
                    int tMs = task.snapshot_times_ms[nextSnapshotIdx];
                    std::string out = appendSuffixBeforeExt(outBase, "_t" + std::to_string(tMs) + "ms");
                    LOG("[Task] snapshot @" + std::to_string(tMs) + "ms -> " + out);
                    win.screenShotEXR(out);
                    writeTaskMetaJson(out, task, win, warmupFrames, countStartAccum);
                    nextSnapshotIdx++;
                }
                if (nextSnapshotIdx >= task.snapshot_times_ms.size()) {
                    saved = true;
                }
            } else if (warmupDone && snapshotSppMode) {
                while (nextSnapshotSppIdx < task.snapshot_spps.size() &&
                       countedSpp >= task.snapshot_spps[nextSnapshotSppIdx]) {
                    std::string outBase = ensureExrPath(win.currentScene.renderOutput);
                    int s = task.snapshot_spps[nextSnapshotSppIdx];
                    std::string out = appendSuffixBeforeExt(outBase, "_spp" + std::to_string(s));
                    LOG("[Task] snapshot @spp=" + std::to_string(s) + " -> " + out);
                    win.screenShotEXR(out);
                    writeTaskMetaJson(out, task, win, warmupFrames, countStartAccum);
                    nextSnapshotSppIdx++;
                }
                if (nextSnapshotSppIdx >= task.snapshot_spps.size()) {
                    saved = true;
                }
            } else if (!saved && warmupDone && countedSpp >= win.currentScene.spp) {
                std::string out = ensureExrPath(win.currentScene.renderOutput);
                LOG("[Task] counted_spp=" + std::to_string(countedSpp) + " reached, saving -> " + out);
                win.screenShotEXR(out);
                writeTaskMetaJson(out, task, win, warmupFrames, countStartAccum);
                saved = true;
            }
            // batch_mode: close once saved; interactive: always keep open
            return !(task.batch_mode && saved);
        });
    } else {
        win.showAndRun();
    }

    if (win.taskTimingInitialized) {
        cudaEventDestroy(win.taskEvRenderStart);
        cudaEventDestroy(win.taskEvRenderEnd);
        win.taskTimingInitialized = false;
    }

    return 0;
}