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

package integrate

import (
	"math"

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

// Adaptive cubature over hyperrectangles: the many-dimensional
// twin of the adaptive Gauss-Legendre quadrature. Each box is measured
// by two product rules (orders 3 and 5 per axis); the difference is
// that box's error estimate, and the globally adaptive loop always
// bisects the worst box along its longest edge, so effort concentrates
// where the integrand actually varies. The 1-D case degenerates to the
// quadrature the package already ships, which doubles as its oracle.

// CubatureOptions tunes IntegrateND. Tolerance bounds the global sum
// of box error estimates (default 1e-10); MaxEvals bounds the function
// evaluations (default two million), an exhausted budget being an
// error naming the achieved estimate, never a silent answer.
type CubatureOptions struct {
	Tolerance float64
	MaxEvals  int
}

// cubBox is one hyperrectangle of the adaptive subdivision: its
// bounds, its measured value and error estimate, and seq, the order
// in which it entered the subdivision. The sequence is the heap's
// tie-break: among equal estimates the earliest inserted box leaves
// first, the same one a scan over the insertion order picks.
type cubBox struct {
	lo, hi []float64
	val    float64
	est    float64
	seq    int
}

// cubBoxAbove reports whether a leaves the box heap before b: the
// larger error estimate first, and among equal estimates the earlier
// insertion. Popping that maximum reproduces the selection of a
// linear scan over the insertion order exactly, ties included, for
// every finite estimate. A non-finite estimate can only come out of
// an overflowed measure, a state in which the integral is already
// meaningless: such a box is a total-order special case and stays at
// the bottom of the heap, leaving every finite estimate to run first,
// where the scan would have left it wherever its insertion happened
// to place it. Two non-finite boxes keep insertion order between
// themselves.
func cubBoxAbove(a, b *cubBox) bool {
	aOut := math.IsNaN(a.est) || math.IsInf(a.est, 0)
	bOut := math.IsNaN(b.est) || math.IsInf(b.est, 0)
	if aOut != bOut {
		return !aOut
	}
	if aOut {
		return a.seq < b.seq
	}
	if a.est != b.est {
		return a.est > b.est
	}
	return a.seq < b.seq
}

// cubSiftUp restores the max-heap order after a push at the tail.
func cubSiftUp(h []*cubBox) {
	i := len(h) - 1
	for i > 0 {
		parent := (i - 1) / 2
		if !cubBoxAbove(h[i], h[parent]) {
			return
		}
		h[i], h[parent] = h[parent], h[i]
		i = parent
	}
}

// cubSiftDown restores the max-heap order after the top has been
// replaced from the tail.
func cubSiftDown(h []*cubBox) {
	n := len(h)
	i := 0
	for {
		left := 2*i + 1
		if left >= n {
			return
		}
		above := left
		if right := left + 1; right < n && cubBoxAbove(h[right], h[left]) {
			above = right
		}
		if !cubBoxAbove(h[above], h[i]) {
			return
		}
		h[i], h[above] = h[above], h[i]
		i = above
	}
}

// cubBoxChunk and cubBoundsChunk size the bisection arenas: one
// allocation per chunk of boxes or bound coordinates instead of one
// per box, so a subdivision that reaches thousands of boxes spends
// tens of allocations, not six per bisection. A chunk never moves once
// handed out, so the heap's pointers stay valid across growth.
const (
	cubBoxChunk    = 256  // boxes per arena chunk
	cubBoundsChunk = 1024 // float64 coordinates per arena chunk
)

// cubBoxArena hands out frozen cubBox values in fixed chunks.
type cubBoxArena struct {
	chunks [][]cubBox
}

// alloc returns the next box, zeroed. A box's fields are written once
// by the caller and never after, which is what lets the heap hold the
// pointer for the life of the subdivision.
func (a *cubBoxArena) alloc() *cubBox {
	if len(a.chunks) == 0 || len(a.chunks[len(a.chunks)-1]) == cubBoxChunk {
		a.chunks = append(a.chunks, make([]cubBox, 0, cubBoxChunk))
	}
	last := len(a.chunks) - 1
	c := append(a.chunks[last], cubBox{})
	a.chunks[last] = c
	return &c[len(c)-1]
}

// cubBoundsArena hands out box-coordinate slices copied from a parent
// box in fixed chunks. A handed-out slice is written once (the copy,
// then the bisected face) and read-only afterwards.
type cubBoundsArena struct {
	chunks [][]float64
	used   int
}

// copy returns src's values in a fresh arena slice.
func (a *cubBoundsArena) copy(src []float64) []float64 {
	n := len(src)
	size := max(n, cubBoundsChunk)
	if len(a.chunks) == 0 || a.used+n > cap(a.chunks[len(a.chunks)-1]) {
		a.chunks = append(a.chunks, make([]float64, 0, size))
		a.used = 0
	}
	last := len(a.chunks) - 1
	c := a.chunks[last]
	keep := len(c)
	c = append(c, src...)
	a.chunks[last] = c
	a.used += n
	return c[keep : keep+n]
}

// IntegrateND returns the integral of f over the hyperrectangle
// [lower, upper] element-wise, by globally adaptive bisection with
// product Gauss-Legendre rules. f receives the evaluation point and
// must not mutate it. A non-finite value, mismatched or empty bounds,
// a reversed edge, or an exhausted evaluation budget is an error.
func IntegrateND(f func(x []float64) float64, lower, upper []float64, opts CubatureOptions) (float64, error) {
	const name = "IntegrateND"
	if len(lower) == 0 || len(lower) != len(upper) {
		return 0, base.Errf("%s: lower and upper must be equal-length non-empty bounds", name)
	}
	for d := range lower {
		if !(upper[d] > lower[d]) {
			return 0, base.Errf("%s: edge %d runs from %g to %g", name, d, lower[d], upper[d])
		}
	}
	tol := opts.Tolerance
	if tol <= 0 {
		tol = 1e-10
	}
	maxEvals := opts.MaxEvals
	if maxEvals <= 0 {
		maxEvals = 2_000_000
	}
	d := len(lower)
	n5, w5, err := GaussLegendreNodes(5)
	if err != nil {
		return 0, base.Errf("%s: %w", name, err)
	}
	n3, w3, err := GaussLegendreNodes(3)
	if err != nil {
		return 0, base.Errf("%s: %w", name, err)
	}

	evals := 0
	// One odometer and one evaluation point serve every rule call:
	// measure is sequential, so each call overwrites what the last
	// read. The bisection budget term is fixed by the dimension.
	idx := make([]int, d)
	point := make([]float64, d)
	// The root box alone costs 5^d + 3^d evaluations before the first
	// budget check could fire, and a bisection calls measure twice,
	// costing 2·(5^d + 3^d); the loop accounts that true cost below.
	// The pre-loop guard is deliberately conservative: it compares
	// against the wider bound 8^d + 6^d, saturating, because a MaxInt
	// budget must not admit a dimension whose true cost merely fits
	// the integer range while needing years to evaluate.
	c5, c3, c8, c6 := 1, 1, 1, 1
	for range d {
		c5 = satMul(c5, 5)
		c3 = satMul(c3, 3)
		c8 = satMul(c8, 8)
		c6 = satMul(c6, 6)
		// A saturated product means the true power left the int range:
		// it is above every budget, and letting it through would put a
		// wrapped count into the later comparisons.
		if c8 == math.MaxInt || c6 == math.MaxInt || c8 > maxEvals || c6 > maxEvals {
			return 0, base.Errf("%s: dimension %d needs more than the %d-evaluation budget for a single bisection", name, d, maxEvals)
		}
	}
	if c5+c3 > maxEvals {
		return 0, base.Errf("%s: dimension %d needs %d evaluations for the root box alone, above the %d budget", name, d, c5+c3, maxEvals)
	}
	boxEvals := 2*c5 + 2*c3
	measure := func(lo, hi []float64) (val, est float64, err error) {
		prodRule := func(nodes, weights []float64) (float64, error) {
			// One odometer over the per-axis nodes; the axis weights
			// multiply along the way, the Jacobian at the end.
			jac := 1.0
			for a := range d {
				jac *= (hi[a] - lo[a]) / 2
			}
			var sum float64
			for {
				for a := range d {
					point[a] = 0.5*(hi[a]-lo[a])*nodes[idx[a]] + 0.5*(hi[a]+lo[a])
				}
				w := jac
				for a := range d {
					w *= weights[idx[a]]
				}
				v := f(point)
				evals++
				if math.IsNaN(v) || math.IsInf(v, 0) {
					return 0, base.Errf("%s: the integrand is non-finite at %v", name, point)
				}
				sum += w * v
				// Odometer advance.
				a := d - 1
				for ; a >= 0; a-- {
					idx[a]++
					if idx[a] < len(nodes) {
						break
					}
					idx[a] = 0
				}
				if a < 0 {
					return sum, nil
				}
			}
		}
		fine, err := prodRule(n5, w5)
		if err != nil {
			return 0, 0, err
		}
		coarse, err := prodRule(n3, w3)
		if err != nil {
			return 0, 0, err
		}
		return fine, math.Abs(fine - coarse), nil
	}

	rootVal, rootEst, err := measure(lower, upper)
	if err != nil {
		return 0, err
	}
	// The boxes awaiting bisection live in a binary max-heap keyed by
	// the error estimate with the insertion sequence as the tie-break,
	// so, while every estimate stays finite, each pop hands back
	// exactly the box a linear scan over the insertion order selects,
	// at logarithmic instead of linear cost; an overflowed measure's
	// non-finite estimate sorts below every finite one. The sifts work
	// index-wise and the backing array grows amortised. The boxes and
	// their bound slices come from the chunk arenas above, one
	// allocation per chunk instead of per box; the root box aliases the
	// caller's bounds, which the solve only reads.
	var boxArena cubBoxArena
	var boundArena cubBoundsArena
	root := boxArena.alloc()
	*root = cubBox{lower, upper, rootVal, rootEst, 0}
	boxes := []*cubBox{root}
	seq := 1
	total := rootVal
	totalEst := rootEst
	// The stopping rule scales the tolerance with the magnitude of the
	// integral, the way IntegrateFunction combines its bounds: the error
	// estimate of an integral of size 1e6 cannot fall below the
	// rounding floor of the sum itself, so a purely absolute tolerance
	// would burn the whole budget and report an exhausted budget
	// instead of the answer. Below unit magnitude the rule is exactly
	// the absolute one it always was.
	for totalEst > tol*math.Max(1, math.Abs(total)) {
		if evals+boxEvals > maxEvals {
			return 0, base.Errf("%s: evaluation budget exhausted (%d), estimate %.6g ± %.2g",
				name, maxEvals, total, totalEst)
		}
		// The heap top is the box with the largest error estimate,
		// the earliest inserted among equals.
		worst := boxes[0]
		if worst.est == 0 {
			break // every box is already exact by the estimate
		}
		// Pop it: move the tail box to the top and sift it down.
		last := len(boxes) - 1
		boxes[0] = boxes[last]
		boxes[last] = nil
		boxes = boxes[:last]
		cubSiftDown(boxes)
		// Bisect along the longest edge.
		longest := 0
		for a := 1; a < d; a++ {
			if worst.hi[a]-worst.lo[a] > worst.hi[longest]-worst.lo[longest] {
				longest = a
			}
		}
		// The dividing plane keeps every other edge: each child is the
		// parent with one face moved to the midpoint, not a corner
		// slice (which would collapse the untouched axes). The
		// unmodified faces stay the parent's own slices, aliased
		// read-only, and the moved face lives in a fresh arena slice:
		// box bounds are never written after their one construction
		// write, so the aliases hold for the life of the heap.
		m := 0.5 * (worst.lo[longest] + worst.hi[longest])
		hi1 := boundArena.copy(worst.hi)
		hi1[longest] = m
		lo2 := boundArena.copy(worst.lo)
		lo2[longest] = m
		v1, e1, err := measure(worst.lo, hi1)
		if err != nil {
			return 0, err
		}
		v2, e2, err := measure(lo2, worst.hi)
		if err != nil {
			return 0, err
		}
		b1 := boxArena.alloc()
		*b1 = cubBox{worst.lo, hi1, v1, e1, seq}
		boxes = append(boxes, b1)
		cubSiftUp(boxes)
		seq++
		b2 := boxArena.alloc()
		*b2 = cubBox{lo2, worst.hi, v2, e2, seq}
		boxes = append(boxes, b2)
		cubSiftUp(boxes)
		seq++
		total += v1 + v2 - worst.val
		totalEst += e1 + e2 - worst.est
	}
	return total, nil
}

// satMul multiplies with saturation at MaxInt, so a power that
// outgrows the int range reads as "above every budget" instead of
// wrapping into a count the comparisons would read as small.
func satMul(a, b int) int {
	if a > math.MaxInt/b {
		return math.MaxInt
	}
	return a * b
}