// pbrt is Copyright(c) 1998-2020 Matt Pharr, Wenzel Jakob, and Greg Humphreys.
// The pbrt source code is licensed under the Apache License, Version 2.0.
// SPDX: Apache-2.0

#include <pbrt/pbrt.h>

#include <pbrt/samplers.h>
#include <pbrt/wavefront/integrator.h>

#include <type_traits>

#ifdef interface
#undef interface
#endif  // interface

namespace pbrt {

// WavefrontPathIntegrator Sampler Methods
void WavefrontPathIntegrator::GenerateRaySamples(int wavefrontDepth, int sampleIndex) {
    auto generateSamples = [=](auto sampler) {
        using ConcreteSampler = std::remove_reference_t<decltype(*sampler)>;
        if constexpr (!std::is_same_v<ConcreteSampler, MLTSampler> &&
                      !std::is_same_v<ConcreteSampler, DebugMLTSampler>)
            GenerateRaySamples<ConcreteSampler>(wavefrontDepth, sampleIndex);
    };
    sampler.DispatchCPU(generateSamples);
}

// template <typename ConcreteSampler>
// void WavefrontPathIntegrator::GenerateRaySamples(int wavefrontDepth, int sampleIndex) {
//     // Generate description string _desc_ for ray sample generation
//     std::string desc = std::string("Generate ray samples - ") + ConcreteSampler::Name();
//
//     RayQueue *rayQueue = CurrentRayQueue(wavefrontDepth);
//     ForAllQueued(
//         desc.c_str(), rayQueue, maxQueueSize, PBRT_CPU_GPU_LAMBDA(const RayWorkItem w) {
//             // Generate samples for ray segment at current sample index
//             // Find first sample dimension
//             int dimension = 6 + 7 * w.depth;
//             if (haveSubsurface)
//                 dimension += 3 * w.depth;
//
//             // Initialize _Sampler_ for pixel, sample index, and dimension
//             ConcreteSampler pixelSampler = *sampler.Cast<ConcreteSampler>();
//             Point2i pPixel = pixelSampleState.pPixel[w.pixelIndex];
//             pixelSampler.StartPixelSample(pPixel, sampleIndex, dimension);
//
//             // Initialize _RaySamples_ structure with sample values
//             RaySamples rs;
//             rs.direct.uc = pixelSampler.Get1D();
//             rs.direct.u = pixelSampler.Get2D();
//             // Initialize remaining samples in _rs_
//             rs.indirect.uc = pixelSampler.Get1D();
//             rs.indirect.u = pixelSampler.Get2D();
//             rs.indirect.rr = pixelSampler.Get1D();
//             // Possibly initialize subsurface subsurface samples in _rs_
//             rs.haveSubsurface = haveSubsurface;
//             if (haveSubsurface) {
//                 rs.subsurface.uc = pixelSampler.Get1D();
//                 rs.subsurface.u = pixelSampler.Get2D();
//             }
//
//             // Store _RaySamples_ in pixel sample state
//             pixelSampleState.samples[w.pixelIndex] = rs;
//         });
// }

template <typename T>
struct is_nile_sampler : std::false_type {};

template <typename Element>
struct is_nile_sampler<NileSampler<Element>> : std::true_type {};

template <typename T>
inline constexpr bool is_nile_sampler_v = is_nile_sampler<T>::value;

template <typename ConcreteSampler>
void WavefrontPathIntegrator::GenerateRaySamples(int wavefrontDepth, int sampleIndex) {
    // Description for queue debug/profiling
    std::string desc = std::string("Generate ray samples - ") + ConcreteSampler::Name();

    RayQueue *rayQueue = CurrentRayQueue(wavefrontDepth);
    ForAllQueued(
        desc.c_str(), rayQueue, maxQueueSize,
        PBRT_CPU_GPU_LAMBDA(const RayWorkItem w) {

            // Concrete sampler instance for this work item (PBRT wavefront style)
            ConcreteSampler pixelSampler = *sampler.Cast<ConcreteSampler>();
            Point2i pPixel = pixelSampleState.pPixel[w.pixelIndex];

            // Compute starting dimension (all branches are compile-time via if constexpr)
            int dimension;
            // -------------------------------------------------------------
            // First change: dimension indexing is more compact
            // -------------------------------------------------------------
            if constexpr (is_nile_sampler_v<ConcreteSampler>) {
                dimension = 4 + (haveSubsurface ? 8 : 6) * w.depth;
            }
            else {
                dimension = 6 + (haveSubsurface ? 10 : 7) * w.depth;
            }

            pixelSampler.StartPixelSample(pPixel, sampleIndex, dimension);

            // One RaySamples object; fill per-branch
            RaySamples rs;
            rs.haveSubsurface = haveSubsurface;

            if constexpr (is_nile_sampler_v<ConcreteSampler>) {
                rs.direct.uc   = pixelSampler.Get1D();
                rs.direct.u    = pixelSampler.Get2D();

                rs.indirect.uc = pixelSampler.Get1D();
                rs.indirect.u  = pixelSampler.Get2D();
                // -------------------------------------------------------------
                // Secondly, we bring this before 1D rr to fit in a 2D pair
                // -------------------------------------------------------------
                if (haveSubsurface) {
                    rs.subsurface.uc = pixelSampler.Get1D();
                    rs.subsurface.u  = pixelSampler.Get2D();
                }
                rs.indirect.rr = pixelSampler.Get1D();

            }
            else {
                // Generic fallback: 1D + 2D like PBRT’s original path
                rs.direct.uc   = pixelSampler.Get1D();
                rs.direct.u    = pixelSampler.Get2D();

                rs.indirect.uc = pixelSampler.Get1D();
                rs.indirect.u  = pixelSampler.Get2D();
                rs.indirect.rr = pixelSampler.Get1D();

                if (haveSubsurface) {
                    rs.subsurface.uc = pixelSampler.Get1D();
                    rs.subsurface.u  = pixelSampler.Get2D();
                }
            }

            // Store the samples
            pixelSampleState.samples[w.pixelIndex] = rs;
        });
}



}  // namespace pbrt
