// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: MIT

package core

import (
	"math/bits"

	"sourcedock.dev/petrbalvin/tensor/internal/engine"
)

// Quasi-random sequences: Halton points, the low-
// discrepancy workhorse of Monte Carlo that needs no tables. Each
// coordinate runs the radical inverse in its own prime base, which
// stratifies every b^k block of points evenly through the unit
// hypercube: the property random sampling only has in expectation.

// HaltonPoints returns the first n Halton points of the given
// dimension as an (n, dim) float64 array, skipping the leading skip
// points (the early Halton coordinates correlate visibly in high
// dimensions; the standard cure is to drop them). dim must be at
// least 1 and at most 32: beyond that the available small primes
// run out and the stratification degrades; a larger request is an
// error, not a silently worse sequence. n+skip must stay below 2^32,
// the index budget both quasi-random constructors enforce; past it
// the point indices wrap the int arithmetic and every point collapses
// back to the origin instead of advancing.
func HaltonPoints(n, dim, skip int) (*Array, error) {
	const name = "HaltonPoints"
	if n < 0 {
		return nil, errf("%s: n must be non-negative, got %d", name, n)
	}
	if dim < 1 || dim > 32 {
		return nil, errf("%s: dim must lie in [1, 32], got %d", name, dim)
	}
	if skip < 0 {
		return nil, errf("%s: skip must be non-negative, got %d", name, skip)
	}
	if uint64(skip)+uint64(n) >= 1<<32 {
		return nil, errf("%s: n + skip must stay below 2^32, the index budget, got %d + %d",
			name, n, skip)
	}
	bases := firstPrimes(dim)
	out := &Array{shape: []int{n, dim}, dt: Float}
	out.alloc(n * dim)
	// The points are independent, so the walk splits over disjoint
	// point ranges: every point writes its own slots and no value is
	// accumulated, so the split cannot move a single coordinate.
	engine.ParallelMin(n, copyMinPerWorker, func(s, e int) {
		for p := s; p < e; p++ {
			idx := p + skip + 1 // the point indexed 0 would be the origin
			for d := range dim {
				out.floats[p*dim+d] = radicalInverse(idx, bases[d])
			}
		}
	})
	return out, nil
}

// radicalInverse reflects the base-b digits of i through the radix
// point: the van der Corput core of every Halton coordinate.
func radicalInverse(i int, b int) float64 {
	f := 1.0
	r := 0.0
	for i > 0 {
		f /= float64(b)
		r += f * float64(i%b)
		i /= b
	}
	return r
}

// firstPrimes returns the first count primes by trial division.
func firstPrimes(count int) []int {
	primes := make([]int, 0, count)
	candidate := 2
	for len(primes) < count {
		isPrime := true
		for _, p := range primes {
			if p*p > candidate {
				break
			}
			if candidate%p == 0 {
				isPrime = false
				break
			}
		}
		if isPrime {
			primes = append(primes, candidate)
		}
		candidate++
	}
	return primes
}

// Sobol sequences: the digital low-discrepancy companion to
// Halton. Each coordinate runs its own linear recurrence over GF(2)
// driven by direction numbers derived from a primitive polynomial, so
// unlike Halton every dimension shares the same base-2 lattice and the
// first 2^k points are exactly stratified through every coordinate,
// not just on average. The parameters below are the Joe and Kuo (2008)
// initialisation table, the current standard, for the first 40
// dimensions; dimension 1 is the plain Gray-coded van der Corput
// sequence and carries no polynomial.

// sobolParams are the initialisation parameters of one dimension: the
// polynomial degree s, the primitive polynomial coefficient a whose
// bits select the recurrence taps, and the s odd initialisation
// integers m_i with 1 <= m_i < 2^i.
type sobolParams struct {
	s uint32
	a uint32
	m []uint32
}

var sobolTable = [...]sobolParams{
	{0, 0, nil},                                     // 1: plain van der Corput
	{1, 0, []uint32{1}},                             // 2
	{2, 1, []uint32{1, 3}},                          // 3
	{3, 1, []uint32{1, 3, 1}},                       // 4
	{3, 2, []uint32{1, 1, 1}},                       // 5
	{4, 1, []uint32{1, 1, 3, 3}},                    // 6
	{4, 4, []uint32{1, 3, 5, 13}},                   // 7
	{5, 2, []uint32{1, 1, 5, 5, 17}},                // 8
	{5, 4, []uint32{1, 1, 5, 5, 5}},                 // 9
	{5, 7, []uint32{1, 1, 7, 11, 19}},               // 10
	{5, 11, []uint32{1, 1, 5, 1, 1}},                // 11
	{5, 13, []uint32{1, 1, 1, 3, 11}},               // 12
	{5, 14, []uint32{1, 3, 5, 5, 31}},               // 13
	{6, 1, []uint32{1, 3, 3, 9, 7, 49}},             // 14
	{6, 13, []uint32{1, 1, 1, 15, 21, 21}},          // 15
	{6, 16, []uint32{1, 3, 1, 13, 27, 49}},          // 16
	{6, 19, []uint32{1, 1, 1, 15, 7, 5}},            // 17
	{6, 22, []uint32{1, 3, 1, 15, 13, 25}},          // 18
	{6, 25, []uint32{1, 1, 5, 5, 19, 61}},           // 19
	{7, 1, []uint32{1, 3, 7, 11, 23, 15, 103}},      // 20
	{7, 4, []uint32{1, 3, 7, 13, 13, 15, 69}},       // 21
	{7, 7, []uint32{1, 1, 3, 13, 7, 35, 63}},        // 22
	{7, 8, []uint32{1, 3, 5, 9, 1, 25, 53}},         // 23
	{7, 14, []uint32{1, 3, 1, 13, 9, 35, 107}},      // 24
	{7, 19, []uint32{1, 3, 1, 5, 27, 61, 31}},       // 25
	{7, 21, []uint32{1, 1, 5, 11, 19, 41, 61}},      // 26
	{7, 28, []uint32{1, 3, 5, 3, 3, 13, 69}},        // 27
	{7, 31, []uint32{1, 1, 7, 13, 1, 19, 1}},        // 28
	{7, 32, []uint32{1, 3, 7, 5, 13, 19, 59}},       // 29
	{7, 37, []uint32{1, 1, 3, 9, 25, 29, 41}},       // 30
	{7, 41, []uint32{1, 3, 5, 13, 23, 1, 55}},       // 31
	{7, 42, []uint32{1, 3, 7, 3, 13, 59, 17}},       // 32
	{7, 50, []uint32{1, 3, 1, 3, 5, 53, 69}},        // 33
	{7, 55, []uint32{1, 1, 5, 5, 23, 33, 13}},       // 34
	{7, 56, []uint32{1, 1, 7, 7, 1, 61, 123}},       // 35
	{7, 59, []uint32{1, 1, 7, 9, 13, 61, 49}},       // 36
	{7, 62, []uint32{1, 3, 3, 5, 3, 55, 33}},        // 37
	{8, 14, []uint32{1, 3, 1, 15, 31, 13, 49, 245}}, // 38
	{8, 21, []uint32{1, 3, 5, 15, 31, 59, 63, 97}},  // 39
	{8, 22, []uint32{1, 3, 1, 11, 11, 11, 77, 249}}, // 40
}

// sobolDirections fills v with the 32-bit direction numbers of the
// given dimension. The first s come straight from the initialisation
// integers scaled into their binary place; the rest follow the
// recurrence v_i = v_{i-s} ^ (v_{i-s} >> s) ^ Σ a_k·v_{i-k} over the
// polynomial taps.
func sobolDirections(p sobolParams, v []uint32) {
	s := int(p.s)
	if s == 0 {
		// Dimension 1: the direction numbers are the binary places
		// themselves, which makes the sequence Gray-coded van der Corput.
		for i := range v {
			v[i] = 1 << (31 - i)
		}
		return
	}
	for i := range s {
		v[i] = p.m[i] << (31 - i)
	}
	for i := s; i < len(v); i++ {
		v[i] = v[i-s] ^ (v[i-s] >> uint(s))
		for k := 1; k < s; k++ {
			if p.a>>(uint(s-1-k))&1 != 0 {
				v[i] ^= v[i-k]
			}
		}
	}
}

// SobolPoints returns the first n Sobol points of the given dimension
// as an (n, dim) float64 array, skipping the leading skip points (the
// sequence starts at the origin index, which carries no information
// and is dropped, exactly as HaltonPoints drops it). dim must be at
// least 1 and at most 40: that is the width of the initialisation
// table, and a larger request is an error, not a silently worse
// sequence. n+skip must stay below 2^32, the period the 32-bit
// direction numbers give; beyond it the index arithmetic would wrap
// back to the origin instead of advancing.
func SobolPoints(n, dim, skip int) (*Array, error) {
	const name = "SobolPoints"
	if n < 0 {
		return nil, errf("%s: n must be non-negative, got %d", name, n)
	}
	if dim < 1 || dim > len(sobolTable) {
		return nil, errf("%s: dim must lie in [1, %d], got %d", name, len(sobolTable), dim)
	}
	if skip < 0 {
		return nil, errf("%s: skip must be non-negative, got %d", name, skip)
	}
	if uint64(skip)+uint64(n) >= 1<<32 {
		return nil, errf("%s: n + skip must stay below 2^32, the sequence period, got %d + %d",
			name, n, skip)
	}
	const width = 32
	out := &Array{shape: []int{n, dim}, dt: Float}
	out.alloc(n * dim)
	// Each coordinate walks its own points. The worker seeds its range's
	// Gray code from scratch, then advances one flip at a time: the Gray
	// codes of consecutive indices differ in the lowest set bit of the
	// newer index, so x picks up exactly the direction numbers the
	// from-scratch walk would XOR, and the exclusive or combines them
	// exactly whatever the order.
	engine.ParallelMin(n, copyMinPerWorker, func(s, e int) {
		for d := range dim {
			var v [width]uint32
			sobolDirections(sobolTable[d], v[:])
			idx := s + skip + 1 // the point indexed 0 would be the origin
			gray := uint32(idx) ^ uint32(idx>>1)
			var x uint32
			for b := range width {
				if gray&(1<<uint(b)) != 0 {
					x ^= v[b]
				}
			}
			for p := s; p < e; p++ {
				out.floats[p*dim+d] = float64(x) / (1 << 32)
				if p+1 < e {
					idx++
					x ^= v[bits.TrailingZeros32(uint32(idx))]
				}
			}
		}
	})
	return out, nil
}