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

package core

import (
	"math"

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

// Elliptic integrals, Jacobi elliptic functions and the Gauss
// hypergeometric function. The complete integrals use the
// arithmetic-geometric mean where it is exact (K directly, E through
// the companion series) and a high-order Gauss-Legendre product where
// it is not (Pi); the Jacobi functions invert the incomplete integral
// F by Newton iteration, which makes sn, cn and dn correct by
// construction: they are the sin, cos and dn of the amplitude whose
// integral is the argument.

// ellipticF evaluates the incomplete integral F(φ, m) = ∫₀^φ dθ / √(1
// − m·sin²θ) through Carlson's symmetric form
//
//	F(φ, m) = sin φ · R_F(cos²φ, 1 − m·sin²φ, 1),
//
// which is exact for every m ∈ [0, 1], φ ∈ [−π/2, π/2]: the integrand's
// endpoint singularity at m tending to 1, φ = π/2 is a zero argument
// of R_F, a regular point of the duplication algorithm. The
// Gauss-Legendre rule this replaces lost accuracy there (2e-16 at
// m = 0.99, 3e-3 at m = 1−1e-6), because a product rule cannot follow a
// square-root singularity.
func ellipticF(phi, m float64) float64 {
	if m == 0 {
		return phi
	}
	// The integrand has period π and a full period integrates to 2K,
	// so any φ reduces to the principal range: F(φ, m) = F(φ̂, m) +
	// 2n·K(m) with φ̂ = φ − nπ the nearest point of [−π/2, π/2].
	n := math.Round(phi / math.Pi)
	phi -= n * math.Pi
	s, c := math.Sincos(phi)
	f := s * carlsonRF(c*c, 1-m*s*s, 1)
	if n != 0 {
		f += 2 * n * EllipticKScalar(m)
	}
	return f
}

// carlsonRF returns Carlson's symmetric elliptic integral of the first
// kind, R_F(x, y, z) = ½∫₀^∞ dt/√((t+x)(t+y)(t+z)), by the
// duplication algorithm (Numerical Recipes, 6.11): each pass averages
// the three arguments, and the series in the final deviations
// converges to double precision in a handful of passes. The arguments
// must be non-negative with at most one zero, which is the case for
// the calls made here.
func carlsonRF(x, y, z float64) float64 {
	// The textbook cut-off is 0.0025, which leaves the third-order
	// series in the deviations good to about 1e-10 relative; halving the
	// deviations quadratically costs two extra passes and buys six
	// digits, so the cut is pushed to 1e-9 here.
	const (
		errtol = 1e-9
		c1     = 1.0 / 24
		c2     = 0.1
		c3     = 3.0 / 44
		c4     = 1.0 / 14
	)
	for range 100 {
		sx, sy, sz := math.Sqrt(x), math.Sqrt(y), math.Sqrt(z)
		lambda := sx*(sy+sz) + sy*sz
		x = 0.25 * (x + lambda)
		y = 0.25 * (y + lambda)
		z = 0.25 * (z + lambda)
		ave := (x + y + z) / 3
		delx := (ave - x) / ave
		dely := (ave - y) / ave
		delz := (ave - z) / ave
		if math.Abs(delx) <= errtol && math.Abs(dely) <= errtol && math.Abs(delz) <= errtol {
			e2 := delx*dely - delz*delz
			e3 := delx * dely * delz
			return (1 + (c1*e2-c2-c3*e3)*e2 + c4*e3) / math.Sqrt(ave)
		}
	}
	// The loop above converges in under ten passes for every admissible
	// input; the fallback keeps the contract of never looping for ever.
	return math.NaN()
}

// EllipticK returns the complete elliptic integral of the first kind
// K(m) = F(π/2, m), evaluated by the arithmetic-geometric mean:
// K(m) = π / (2·AGM(1, √(1−m))). The parameter convention is m = k².
// m must be below 1; as m tends to 1 it diverges
// and K returns +Inf there, m = 1 included as the limit only through
// the caller's rounding.
func EllipticK(m *Array) (*Array, error) {
	return m.realFunc("EllipticK", func(v float64) float64 {
		if math.IsNaN(v) {
			return math.NaN()
		}
		if v >= 1 {
			if v == 1 {
				return math.Inf(1)
			}
			return math.NaN()
		}
		if v < 0 {
			// Negative parameter: transform to a positive one,
			// K(−s) = K(s/(1+s))/√(1+s). The complement 1 − s/(1+s) =
			// 1/(1+s) is handed to the AGM directly: for s above 2^53
			// the float64 sum 1+s rounds to s, so the transformed
			// parameter would round to exactly 1 and the AGM would
			// return its round-off floor (~1.8e15) instead of the true
			// K, which stays finite, and decays to 0, for every finite
			// s.
			s := -v
			return agmKFromComplement(1/(1+s)) / math.Sqrt(1+s)
		}
		return EllipticKScalar(v)
	})
}

// EllipticKScalar is K(m) at one point (0 ≤ m < 1) by the AGM.
func EllipticKScalar(m float64) float64 {
	return agmKFromComplement(1 - m)
}

// agmKFromComplement is K(1−ε) = π/(2·AGM(1, √ε)) evaluated from the
// complementary parameter ε = 1−m, the form the negative-parameter
// transform needs: there ε = 1/(1+s) stays exact where 1−ε would round
// to 1.
func agmKFromComplement(eps float64) float64 {
	a, b := 1.0, math.Sqrt(eps)
	// The stop test sits a few ulps above machine zero: rounding can
	// pin a and b one ulp apart forever, and any threshold below that
	// is an infinite loop, not extra accuracy.
	for math.Abs(a-b) > 4*epsF*math.Max(1, a) {
		a, b = 0.5*(a+b), math.Sqrt(a*b)
	}
	return math.Pi / (2 * a)
}

// EllipticE returns the complete elliptic integral of the second kind
// E(m) = ∫₀^{π/2} √(1 − m·sin²θ) dθ through the AGM companion series
// E = K·(1 − Σ 2^{n−1} c_n²), with c_n² = a_n² − b_n² the AGM
// remainders. m = 1 gives 1; m > 1 is NaN; negative m transforms like
// K's does.
func EllipticE(m *Array) (*Array, error) {
	return m.realFunc("EllipticE", ellipticEScalar)
}

// ellipticEScalar is E(m) at one point.
func ellipticEScalar(m float64) float64 {
	if math.IsNaN(m) {
		return math.NaN()
	}
	if m == 1 {
		return 1
	}
	if m > 1 {
		return math.NaN()
	}
	if m < 0 {
		// E(−s) = √(1+s)·E(s/(1+s)).
		s := -m
		return math.Sqrt(1+s) * ellipticEScalar(s/(1+s))
	}
	a, b := 1.0, math.Sqrt(1-m)
	k := EllipticKScalar(m)
	sum := 0.0
	pow2 := 0.5 // 2^{n-1} starting at n = 1
	for math.Abs(a-b) > 4*epsF*math.Max(1, a) {
		c2 := a*a - b*b
		sum += pow2 * c2
		pow2 *= 2
		a, b = 0.5*(a+b), math.Sqrt(a*b)
	}
	return k * (1 - sum)
}

// EllipticPi returns the complete elliptic integral of the third kind
// Π(n, m) = ∫₀^{π/2} dθ / ((1 − n·sin²θ)√(1 − m·sin²θ)) element-wise
// over paired arrays, through Carlson's symmetric forms
//
//	Π(n, m) = R_F(0, 1−m, 1) + (n/3)·R_J(0, 1−m, 1, 1−n),
//
// whose duplication algorithms are exact at the singular ends of the
// parameter square: full double precision for every n < 1 and m < 1,
// where the 64-point product rule this replaces lost seven digits at
// m = 1−1e-6 and more as n approached 1 (measured against mpmath: the
// worst relative error over the pinned table is 6e-16). At n = 1 and
// at m = 1 the integral diverges: the value there is +Inf, and beyond
// it NaN.
func EllipticPi(n, m *Array) (*Array, error) {
	if !sameShape(n.shape, m.shape) {
		return nil, errf("EllipticPi: shape mismatch %s vs %s", shapeText(n.shape), shapeText(m.shape))
	}
	if n.dt == Complex || m.dt == Complex {
		return nil, errf("EllipticPi: complex arrays are not supported")
	}
	out := &Array{shape: append([]int{}, n.shape...), dt: Float}
	out.alloc(n.Len())
	engine.Parallel(n.Len(), func(s, e int) {
		for i := s; i < e; i++ {
			out.floats[i] = ellipticPiScalar(n.floatAt(i), m.floatAt(i))
		}
	})
	return out, nil
}

// ellipticPiScalar is Π(n, m) at one point, through Carlson's
// symmetric forms:
//
//	Π(n, m) = R_F(0, 1−m, 1) + (n/3)·R_J(0, 1−m, 1, 1−n).
//
// Both forms are exact at the singular ends of the parameter square,
// which is where the product rule this replaces lost its digits.
func ellipticPiScalar(n, m float64) float64 {
	if math.IsNaN(n) || math.IsNaN(m) || m >= 1 || n >= 1 {
		if m == 1 || n == 1 {
			return math.Inf(1)
		}
		return math.NaN()
	}
	rf := carlsonRF(0, 1-m, 1)
	if n == 0 {
		return rf
	}
	return rf + n/3*carlsonRJ(0, 1-m, 1, 1-n)
}

// carlsonRC returns Carlson's degenerate symmetric integral
// R_C(x, y) = R_F(x, y, y) for x ≥ 0 and y ≠ 0, by the duplication
// algorithm with the single-deviation series (Carlson 1995,
// "Numerical computation of real or complex elliptic integrals",
// (16)-(20)). A negative y is its Cauchy principal value, through the
// paper's (21); that case never arises for the calls made here, and is
// implemented so the function stands on its own.
func carlsonRC(x, y float64) float64 {
	// r is the target relative truncation error; 1e-16 asks for double
	// precision, and the seven-term series is good to that.
	const r = 1e-16
	if y < 0 {
		// (21) with the positive magnitude: R_C(x, −Y) =
		// sqrt(x/(x+Y))·R_C(x+Y, Y).
		return math.Sqrt(x/(x-y)) * carlsonRC(x-y, -y)
	}
	a0 := (x + 2*y) / 3
	q := math.Pow(3*r, -0.125) * math.Abs(a0-x)
	a, xc, yc := a0, x, y
	pow4 := 1.0 // 4^{−m}
	for range 200 {
		sx, sy := math.Sqrt(xc), math.Sqrt(yc)
		lambda := 2*sx*sy + yc
		an := (a + lambda) / 4
		pow4n := pow4 / 4
		if pow4n*q < math.Abs(an) {
			s := (y - a0) / (an / pow4n)
			s2 := s * s
			series := 1 + (3.0/10)*s2 + (1.0/7)*s2*s + (3.0/8)*s2*s2 +
				(9.0/22)*s2*s2*s + (159.0/208)*s2*s2*s2 + (9.0/8)*s2*s2*s2*s
			return series / math.Sqrt(an)
		}
		xc, yc = (xc+lambda)/4, (yc+lambda)/4
		a, pow4 = an, pow4n
	}
	return math.NaN()
}

// carlsonRJ returns Carlson's symmetric integral
// R_J(x, y, z, p) = (3/2)∫₀^∞ dt/((t+p)√((t+x)(t+y)(t+z))) for x, y,
// z ≥ 0 with at most one zero and p > 0, by the duplication theorem
// with the five-variable series (Carlson 1995, (24)-(32)): each pass
// averages the variables through lambda and accumulates the
// R_C(1, 1+e) correction the theorem contributes, then a sixth-order
// series in the deviations from the mean finishes the job.
func carlsonRJ(x, y, z, p float64) float64 {
	const r = 1e-16
	if p <= 0 {
		// The principal value for negative p is the paper's (33),
		// which needs a permutation of x, y, z; none of the callers
		// reaches it, so it is refused rather than approximated.
		return math.NaN()
	}
	a0 := (x + y + z + 2*p) / 5
	delta := (p - x) * (p - y) * (p - z)
	q := math.Pow(r/4, -1.0/6) * math.Max(math.Max(math.Abs(a0-x), math.Abs(a0-y)),
		math.Max(math.Abs(a0-z), math.Abs(a0-p)))
	a, xc, yc, zc, pc := a0, x, y, z, p
	pow4 := 1.0
	total := 0.0
	for range 200 {
		sx, sy, sz, sp := math.Sqrt(xc), math.Sqrt(yc), math.Sqrt(zc), math.Sqrt(pc)
		lambda := sx*sy + sx*sz + sy*sz
		an := (a + lambda) / 4
		pow4n := pow4 / 4
		// The m-th correction carries 4^{−3m}; pow4 is 4^{−m} already.
		d := (sp + sx) * (sp + sy) * (sp + sz)
		e := pow4 * pow4 * pow4 * delta / (d * d)
		total += 6 * pow4 / d * carlsonRC(1, 1+e)
		if pow4n*q < math.Abs(an) {
			scaled := an / pow4n // 4^n·A_n
			xx := (a0 - x) / scaled
			yy := (a0 - y) / scaled
			zz := (a0 - z) / scaled
			pp := (-xx - yy - zz) / 2
			e2 := xx*yy + xx*zz + yy*zz - 3*pp*pp
			e3 := xx*yy*zz + 2*e2*pp + 4*pp*pp*pp
			e4 := (2*xx*yy*zz + e2*pp + 3*pp*pp*pp) * pp
			e5 := xx * yy * zz * pp * pp
			series := 1 - (3.0/14)*e2 + (1.0/6)*e3 + (9.0/88)*e2*e2 -
				(3.0/22)*e4 - (9.0/52)*e2*e3 + (3.0/26)*e5
			return pow4n*series/(an*math.Sqrt(an)) + total
		}
		xc, yc, zc, pc = (xc+lambda)/4, (yc+lambda)/4, (zc+lambda)/4, (pc+lambda)/4
		a, pow4 = an, pow4n
	}
	return math.NaN()
}

// EllipticFScalar is the incomplete elliptic integral of the first
// kind F(φ, m) = ∫₀^φ dθ/√(1−m·sin²θ) at one point, through Carlson's
// symmetric form; valid for every m ∈ [0, 1) and any real φ. The
// inverse view is the amplitude: u = F(φ, m) means φ = am(u, m), the
// reading behind the Jacobi functions.
func EllipticFScalar(phi, m float64) float64 {
	return ellipticF(phi, m)
}

// JacobiSN returns sn(u, m), the Jacobi elliptic sine, element-wise
// over u with the parameter m shared: sn inverts the incomplete
// integral u = F(φ, m) through φ, giving sn = sin φ. The same
// inversion serves cn and dn, which are the cos and the sqrt factor
// of the same amplitude. Valid for m ∈ [0, 1) and any real u, where
// m = 0 degenerates to the circular functions; m = 1 is refused
// because the amplitude would have to travel through the singular
// complete integral.
func JacobiSN(u *Array, m float64) (*Array, error) {
	return jacobi(u, m, true, func(phi, mm float64) float64 { return math.Sin(phi) })
}

// JacobiCN returns cn(u, m) element-wise.
func JacobiCN(u *Array, m float64) (*Array, error) {
	return jacobi(u, m, false, func(phi, mm float64) float64 { return math.Cos(phi) })
}

// JacobiDN returns dn(u, m) element-wise.
func JacobiDN(u *Array, m float64) (*Array, error) {
	return jacobi(u, m, false, func(phi, mm float64) float64 {
		return math.Sqrt(1 - mm*math.Sin(phi)*math.Sin(phi))
	})
}

// jacobi inverts u = F(φ, m) by a bracketed Newton iteration and maps
// the amplitude through f.
func jacobi(u *Array, m float64, odd bool, f func(phi, mm float64) float64) (*Array, error) {
	if u.dt == Complex {
		return nil, errf("Jacobi: complex arrays are not supported")
	}
	if math.IsNaN(m) || m < 0 || m >= 1 {
		return nil, errf("Jacobi: the parameter m must lie in [0, 1), got %g", m)
	}
	// K scales the period; the amplitude search brackets φ in [0, hi]
	// by doubling until F covers u.
	k := EllipticKScalar(m)
	out := &Array{shape: append([]int{}, u.shape...), dt: Float}
	out.alloc(u.Len())
	engine.Parallel(u.Len(), func(s, e int) {
		for i := s; i < e; i++ {
			uu := u.floatAt(i)
			// Sign symmetry first: work with |u|.
			sign := 1.0
			if uu < 0 {
				sign = -1
				uu = -uu
			}
			// Bracket: F(φ) ≥ φ·(1−m)^... the integrand is at least
			// 1 on [0, π/2] and periodic beyond; hi = u + K covers
			// every u ≥ 0 because F(u + K-margin)... doubling is the
			// safe route.
			hi := math.Max(uu, k)
			for ellipticF(hi, m) < uu {
				hi *= 2
			}
			lo := 0.0
			phi := 0.5 * (lo + hi)
			for range 200 {
				// Newton from the midpoint of the current bracket,
				// re-bracketed every step: F is strictly increasing.
				val := ellipticF(phi, m)
				if val < uu {
					lo = phi
				} else {
					hi = phi
				}
				den := math.Sqrt(1 - m*math.Sin(phi)*math.Sin(phi))
				// F'(φ) = 1/den, so the Newton step multiplies by den.
				next := phi + (uu-val)*den
				if next <= lo || next >= hi {
					next = 0.5 * (lo + hi)
				}
				if math.Abs(next-phi) < 1e-16*math.Max(1, math.Abs(phi)) {
					phi = next
					break
				}
				phi = next
			}
			if odd {
				out.floats[i] = sign * f(phi, m)
			} else {
				out.floats[i] = f(phi, m) // dn is an even function
			}
		}
	})
	return out, nil
}

// JacobiCDScalar is cd(u, m) = cn(u, m)/dn(u, m) at one point, by the
// Gauss AGM: the descending Landen sequence converges quadratically
// and the amplitude folds back down the chain, so a handful of passes
// covers any u. The parameter convention is m = k². Valid for m ∈ [0,
// 1) and any real u, where m = 0 degenerates cd to cos. This is the
// function the elliptic rational function needs for the zeroes of a
// Cauer filter design.
func JacobiCDScalar(u, m float64) float64 {
	if math.IsNaN(m) || m < 0 || m >= 1 || math.IsNaN(u) {
		return math.NaN()
	}
	if m == 0 {
		return math.Cos(u)
	}
	// The AGM chain: a ascends to the mean, b descends, c is half
	// their gap and vanishes quadratically.
	k := EllipticKScalar(m)
	// cd has period 4K and antisymmetry about 2K: reduce to [0, 2K)
	// and carry the sign.
	sign := 1.0
	u = math.Mod(u, 4*k)
	if u < 0 {
		u += 4 * k
	}
	if u >= 2*k {
		u -= 2 * k
		sign = -1
	}
	var a, b, c [64]float64
	a[0], b[0], c[0] = 1, math.Sqrt(1-m), math.Sqrt(m)
	n := 0
	for n < 63 {
		n++
		a[n] = 0.5 * (a[n-1] + b[n-1])
		b[n] = math.Sqrt(a[n-1] * b[n-1])
		c[n] = 0.5 * (a[n-1] - b[n-1])
		// The gap between a and b stops shrinking a few ulps above
		// machine zero, so the stop test carries the same floor the
		// complete integral's AGM uses; below it more passes would be
		// an infinite loop, not extra accuracy.
		if math.Abs(c[n]) <= 4*epsF*math.Abs(a[n]) {
			break
		}
	}
	// The amplitude climbs the chain, then folds back down it.
	phi := float64(int(1)<<uint(n)) * a[n] * u
	for i := n; i > 0; i-- {
		phi = 0.5 * (phi + math.Asin(c[i]/a[i]*math.Sin(phi)))
	}
	return sign * math.Cos(phi) / math.Sqrt(1-m*math.Sin(phi)*math.Sin(phi))
}

// Hypergeometric2F1 returns the Gauss hypergeometric function
// ₂F₁(a, b; c; x) = Σ (a)_k (b)_k / (c)_k · x^k / k! element-wise
// over x, with a, b, c scalar parameters. The series converges for
// |x| < 1; a or b a non-positive integer terminates it as a
// polynomial. A non-positive integer c is an error. Outside the disc
// of convergence the result is the Gauss value at x = 1 when c > a+b
// (finite), the divergence limit +Inf when x = 1 and c ≤ a+b, and NaN
// for every other |x| ≥ 1.
func Hypergeometric2F1(a, b, c float64, x *Array) (*Array, error) {
	if c == math.Trunc(c) && c <= 0 {
		return nil, errf("Hypergeometric2F1: c must not be a non-positive integer, got %g", c)
	}
	terminated := (a == math.Trunc(a) && a <= 0) || (b == math.Trunc(b) && b <= 0)
	return x.realFunc("Hypergeometric2F1", func(v float64) float64 {
		if math.IsNaN(v) {
			return math.NaN()
		}
		if math.Abs(v) >= 1 && !terminated {
			if v == 1 {
				// The Gauss value at x = 1 exists when c > a+b.
				if c > a+b {
					return math.Gamma(c) * math.Gamma(c-a-b) / (math.Gamma(c-a) * math.Gamma(c-b))
				}
				return math.Inf(1)
			}
			return math.NaN()
		}
		term := 1.0
		sum := 1.0
		for k := 1; k <= 100000; k++ {
			term *= (a + float64(k-1)) * (b + float64(k-1)) / (c + float64(k-1)) * v / float64(k)
			sum += term
			if math.Abs(term) < 1e-18*math.Abs(sum) {
				break
			}
			if math.IsInf(sum, 0) {
				return sum
			}
		}
		return sum
	})
}