#include "Algorithm/Sobol.h"

//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);
    }

    // Dimensions 2 to 6
    for (int d = 1; d < SOBOL_DIMENSIONS; d++)
    {
        int s = s_vals[d - 1];
        int a = a_vals[d - 1];
        const int* m = m_vals[d - 1];

        // Set initial direction numbers
        for (int i = 0; i < s; i++)
        {
            sobol_direction_vectors[d][i] = (uint32_t)m[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;
        }
    }
}

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

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

float sobol_sample(uint32_t index, uint32_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_next(sobol_state_t* state)
{
    uint32_t index = state->index;
    uint32_t dimension = state->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];
        }
    }

    // Increment the index for the next call
    state->dimension = (state->dimension + 1) % SOBOL_DIMENSIONS;

    return sobol_uint_to_float(result);
}