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

package integrate

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

import (
	"math"
	"sync"
)

// Numerical quadrature: the definite integral of a function over an
// interval. The library integrates ODEs but until now not plain
// integrals, which every data-reduction pipeline needs.
//
// The scheme is adaptive Gauss-Legendre with an honest internal
// error estimate: every subinterval is evaluated by a 21-point rule
// and a 10-point rule over the same interval, the gap between the two
// is that subinterval's error, and the subinterval with the largest
// error is bisected until the summed error meets the tolerance. The
// 21 and 10 point nodes come from Newton's method on the Legendre
// recurrence, so the whole construction is derived in the library
// rather than imported as a table. Infinite bounds map onto (0, 1) by
// rational substitution before the rule runs; an integrand over an
// infinite interval has to decay to zero for this to converge.

// QuadratureOptions tunes IntegrateFunction. RelTol ≤ 0 means 1e-10,
// AbsTol ≤ 0 means 1e-12, MaxIntervals ≤ 0 means 256.
type QuadratureOptions struct {
	RelTol       float64
	AbsTol       float64
	MaxIntervals int
}

// gaussLegendreEntry holds the cached nodes and weights of the
// n-point Gauss-Legendre rule over [-1, 1].
type gaussLegendreEntry struct {
	nodes   []float64
	weights []float64
}

var (
	gaussLegendreCacheMu sync.Mutex
	gaussLegendreCache   = map[int]gaussLegendreEntry{}
)

// GaussLegendreNodes returns the nodes and weights of the n-point
// Gauss-Legendre rule over [-1, 1], exact for polynomials up to
// degree 2n−1. Nodes come out ascending. n must be between 1 and 128.
// The nodes are the roots of the n-th Legendre polynomial, found by
// Newton's method on the three-term recurrence, which reaches
// rounding-level accuracy in a handful of iterations per node. The
// returned slices are a shared cache: they must be treated as
// read-only, because a write would poison every later quadrature run
// on the same node count.
func GaussLegendreNodes(n int) (nodes, weights []float64, err error) {
	if n < 1 || n > 128 {
		return nil, nil, base.Errf("GaussLegendreNodes: n must be between 1 and 128, got %d", n)
	}
	gaussLegendreCacheMu.Lock()
	entry, ok := gaussLegendreCache[n]
	gaussLegendreCacheMu.Unlock()
	if ok {
		return entry.nodes, entry.weights, nil
	}
	nodes = make([]float64, n)
	weights = make([]float64, n)
	for i := range n {
		// The trigonometric starting guess separates the roots well
		// enough that Newton never jumps to a neighbour.
		x := -math.Cos(math.Pi * (float64(i) + 0.75) / (float64(n) + 0.5))
		var p, dp float64
		for range 100 {
			p, dp = legendrePair(n, x)
			dx := p / dp
			x -= dx
			if math.Abs(dx) <= 1e-16*(1+math.Abs(x)) {
				break
			}
		}
		if math.Abs(p) > 1e-10 {
			return nil, nil, base.Errf("GaussLegendreNodes: Newton failed to converge for n=%d", n)
		}
		nodes[i] = x
		weights[i] = 2 / ((1 - x*x) * dp * dp)
	}
	entry = gaussLegendreEntry{nodes: nodes, weights: weights}
	gaussLegendreCacheMu.Lock()
	gaussLegendreCache[n] = entry
	gaussLegendreCacheMu.Unlock()
	return nodes, weights, nil
}

// legendrePair evaluates P_n(x) and P_n'(x) by the three-term
// recurrence and its derivative identity.
func legendrePair(n int, x float64) (p, dp float64) {
	p = 1
	if n == 0 {
		return 1, 0
	}
	pm := 0.0
	for k := 1; k <= n; k++ {
		pm, p = p, ((2*float64(k)-1)*x*p-float64(k-1)*pm)/float64(k)
	}
	dp = float64(n) * (x*p - pm) / (x*x - 1)
	return p, dp
}

// IntegrateFunction returns the definite integral of f over [a, b] with an
// estimate of the absolute error. Infinite bounds are accepted: a = −∞
// or b = +∞ (or both) integrate over the whole tail under a rational
// substitution, which asks the integrand to decay to zero. A reversed
// interval (a > b) integrates in the negative direction. A tolerance
// that cannot be met within MaxIntervals subintervals is an error,
// never a silent approximation.
//
// Every sampled scheme has one blind spot: a feature entirely inside
// the gaps of the first rule's nodes, say a peak far narrower than
// (b−a)/n, produces small values everywhere it samples and is missed
// with a small error estimate. Split known sharp features into their
// own IntegrateFunction calls.
func IntegrateFunction(f func(x float64) (float64, error), a, b float64, opts QuadratureOptions) (float64, float64, error) {
	if opts.RelTol <= 0 {
		opts.RelTol = 1e-10
	}
	if opts.AbsTol <= 0 {
		opts.AbsTol = 1e-12
	}
	if opts.MaxIntervals <= 0 {
		opts.MaxIntervals = 256
	}
	if math.IsNaN(a) || math.IsNaN(b) {
		return 0, 0, base.Errf("IntegrateFunction: bounds must not be NaN")
	}
	sign := 1.0
	if b < a {
		a, b = b, a
		sign = -1
	}
	if a == b {
		return 0, 0, nil
	}
	// Map every bound combination onto a plain finite interval with a
	// wrapped integrand carrying the substitution's Jacobian.
	lo, hi := 0.0, 1.0
	g := f
	switch {
	case a == math.Inf(-1) && b == math.Inf(1):
		g = func(t float64) (float64, error) {
			u := t - 0.5
			den := 0.25 - u*u
			if den <= 0 {
				return 0, nil
			}
			x := 2 * u / den
			fx, ferr := f(x)
			if ferr != nil {
				return 0, ferr
			}
			return fx * (0.5 + 2*u*u) / (den * den), nil
		}
	case a == math.Inf(-1):
		g = func(t float64) (float64, error) {
			x := b - t/(1-t)
			fx, ferr := f(x)
			if ferr != nil {
				return 0, ferr
			}
			return fx / (1 - t) / (1 - t), nil
		}
	case b == math.Inf(1):
		g = func(t float64) (float64, error) {
			x := a + t/(1-t)
			fx, ferr := f(x)
			if ferr != nil {
				return 0, ferr
			}
			return fx / (1 - t) / (1 - t), nil
		}
	default:
		lo, hi = a, b
	}

	nodes10, w10, err := GaussLegendreNodes(10)
	if err != nil {
		return 0, 0, err
	}
	nodes21, w21, err := GaussLegendreNodes(21)
	if err != nil {
		return 0, 0, err
	}
	rule := func(g func(float64) (float64, error), l, r float64, xs, ws []float64) (float64, error) {
		mid, half := (l+r)/2, (r-l)/2
		total := 0.0
		for i := range xs {
			fx, ferr := g(mid + half*xs[i])
			if ferr != nil {
				return 0, ferr
			}
			total += ws[i] * fx
		}
		return total * half, nil
	}

	// Each leaf carries its 21-point value and its error, the gap to
	// the 10-point rule over the same interval.
	type leaf struct {
		l, r, value, err float64
	}
	measure := func(l, r float64) (leaf, error) {
		v21, err := rule(g, l, r, nodes21, w21)
		if err != nil {
			return leaf{}, err
		}
		v10, err := rule(g, l, r, nodes10, w10)
		if err != nil {
			return leaf{}, err
		}
		// An infinite integrand value would poison errSum with NaN and
		// slip through the loop condition (NaN comparisons are false),
		// publishing a bogus integral, so both non-finite kinds are
		// refused exactly as IntegrateND refuses them.
		if math.IsNaN(v21) || math.IsNaN(v10) || math.IsInf(v21, 0) || math.IsInf(v10, 0) {
			return leaf{}, base.Errf("the integrand returned a non-finite value on [%g, %g]", l, r)
		}
		return leaf{l: l, r: r, value: v21, err: math.Abs(v21 - v10)}, nil
	}

	first, err := measure(lo, hi)
	if err != nil {
		return 0, 0, base.Errf("IntegrateFunction: %w", err)
	}
	leaves := []leaf{first}
	budget := func() (total, errSum float64) {
		for _, s := range leaves {
			total += s.value
			errSum += s.err
		}
		return total, errSum
	}
	value, errSum := budget()
	for errSum > math.Max(opts.AbsTol, opts.RelTol*math.Abs(value)) {
		if len(leaves) >= opts.MaxIntervals {
			return 0, 0, base.Errf("IntegrateFunction: error estimate %g exceeds the tolerance within %d subintervals",
				errSum, opts.MaxIntervals)
		}
		// Bisect the leaf that contributes the most error.
		worst := 0
		for i, s := range leaves {
			if s.err > leaves[worst].err {
				worst = i
			}
		}
		w := leaves[worst]
		left, lerr := measure(w.l, (w.l+w.r)/2)
		if lerr != nil {
			return 0, 0, base.Errf("IntegrateFunction: %w", lerr)
		}
		right, rerr := measure((w.l+w.r)/2, w.r)
		if rerr != nil {
			return 0, 0, base.Errf("IntegrateFunction: %w", rerr)
		}
		leaves[worst] = left
		leaves = append(leaves, right)
		value, errSum = budget()
	}
	return sign * value, errSum, nil
}