#include "Algorithm/Sobol.h"

static inline int sobol_get_bit(uint32_t index, uint32_t bit)
{
    return (index >> bit) & 1;
}

//NOTE: Maybe precompute the vector table?
void sobol_init()
{
    // First, set dimension 0 manually (Van der Corput)
    for (int i = 0; i < SOBOL_BITS; i++)
    {
        sobol_direction_vectors[0][i] = 1U << (31 - i);
    }

    for (int d = 1; d < SOBOL_DIMENSIONS; d++)
    {
        int s = s_vals[d - 1];
        int a = a_vals[d - 1];

        // Set initial direction numbers
        for (int i = 0; i < s; i++)
        {
            sobol_direction_vectors[d][i] = m_vals[d - 1][i] << (31 - i);
        }

        // Compute remaining direction numbers
        for (int i = s; i < SOBOL_BITS; i++)
        {
            uint32_t v = sobol_direction_vectors[d][i - s];
            v ^= v >> s;

            for (int k = 1; k < s; k++)
            {
                if ((a >> (s - 1 - k)) & 1)
                    v ^= sobol_direction_vectors[d][i - k];
            }

            sobol_direction_vectors[d][i] = v;
        }
    }
}

uint16_t sobol_get_dimension(uint16_t bounce, sampling_dimension_t operation_type)
{
    return PRNG_BASE_NUM + bounce * PRNG_BOUNCE_NUM + operation_type;
}

static inline float sobol_uint_to_float(uint32_t x)
{
    return x * 2.3283064365386963e-10f; // 1/2^32
}

float sobol_sample(uint32_t index, uint16_t dimension)
{
    if (dimension >= SOBOL_DIMENSIONS)
    {
        return 0.0f; // Invalid dimension
    }

    uint32_t result = 0;
    for (int i = 0; i < SOBOL_BITS; i++)
    {
        if (sobol_get_bit(index, i))
        {
            result ^= sobol_direction_vectors[dimension][i];
        }
    }

    return sobol_uint_to_float(result);
}

float sobol_sample_scrambled(uint32_t index, uint16_t dimension, uint32_t scramble)
{
    if (dimension >= SOBOL_DIMENSIONS)
    {
        return 0.0f;
    }

    uint32_t result = 0;
    for (int i = 0; i < SOBOL_BITS; i++)
    {
        if (sobol_get_bit(index, i))
        {
            result ^= sobol_direction_vectors[dimension][i];
        }
    }

    // Apply XOR scrambling to decorrelate pixels
    result ^= scramble;

    return sobol_uint_to_float(result);
}