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

package core

import (
	"math"
	"math/big"

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

// Classical orthogonal polynomials on the real line, each element-wise
// over real arrays with integer degree or order. The recurrences are
// the stable ones for every family, and the normalisations follow the
// standard physics conventions: Hermite physicists' Hₙ with leading
// coefficient 2ⁿ, Laguerre Lₙ (α = 0) and generalised Lₙ^α, Chebyshev
// of the first and second kind Tₙ and Uₙ on [−1, 1].

// Hermite returns the physicists' Hermite polynomial Hₙ(x) of each
// element, by the recurrence Hₙ₊₁ = 2x·Hₙ − 2n·Hₙ₋₁.
func Hermite(n int, x *Array) (*Array, error) {
	if n < 0 {
		return nil, errf("Hermite: degree must be ≥ 0, got %d", n)
	}
	if x.dt == Complex {
		return nil, errf("Hermite: complex arrays are not supported")
	}
	out := &Array{shape: append([]int{}, x.shape...), dt: Float}
	out.alloc(x.Len())
	engine.Parallel(x.Len(), func(s, e int) {
		for i := s; i < e; i++ {
			xv := x.floatAt(i)
			h0, h1 := 1.0, 2*xv
			for k := 1; k < n; k++ {
				h0, h1 = h1, 2*xv*h1-2*float64(k)*h0
			}
			if n == 0 {
				h1 = 1
			}
			out.floats[i] = h1
		}
	})
	return out, nil
}

// Laguerre returns the generalised Laguerre polynomial Lₙ^α(x) of each
// element, by the recurrence (n+1)Lₙ₊₁ = (2n+1+α−x)Lₙ − (n+α)Lₙ₋₁.
// The degree n must be ≥ 0.
func Laguerre(n int, alpha float64, x *Array) (*Array, error) {
	if n < 0 {
		return nil, errf("Laguerre: degree must be ≥ 0, got %d", n)
	}
	if x.dt == Complex {
		return nil, errf("Laguerre: complex arrays are not supported")
	}
	if n == 0 {
		out := &Array{shape: append([]int{}, x.shape...), dt: Float}
		out.alloc(x.Len())
		for i := range x.Len() {
			out.floats[i] = 1
		}
		return out, nil
	}
	out := &Array{shape: append([]int{}, x.shape...), dt: Float}
	out.alloc(x.Len())
	engine.Parallel(x.Len(), func(s, e int) {
		for i := s; i < e; i++ {
			xv := x.floatAt(i)
			l0, l1 := 1.0, 1.0+float64(alpha)-xv
			for k := 1; k < n; k++ {
				kf := float64(k)
				l0, l1 = l1, ((2*kf+1+alpha-xv)*l1-(kf+alpha)*l0)/(kf+1)
			}
			out.floats[i] = l1
		}
	})
	return out, nil
}

// ChebyshevT returns the Chebyshev polynomial of the first kind
// Tₙ(x) = cos(n·arccos x) of each element, by the recurrence
// Tₙ₊₁ = 2x·Tₙ − Tₙ₋₁.
func ChebyshevT(n int, x *Array) (*Array, error) {
	if n < 0 {
		return nil, errf("ChebyshevT: degree must be ≥ 0, got %d", n)
	}
	if x.dt == Complex {
		return nil, errf("ChebyshevT: complex arrays are not supported")
	}
	out := &Array{shape: append([]int{}, x.shape...), dt: Float}
	out.alloc(x.Len())
	engine.Parallel(x.Len(), func(s, e int) {
		for i := s; i < e; i++ {
			xv := x.floatAt(i)
			t0, t1 := 1.0, xv
			for k := 1; k < n; k++ {
				t0, t1 = t1, 2*xv*t1-t0
			}
			if n == 0 {
				t1 = 1
			}
			out.floats[i] = t1
		}
	})
	return out, nil
}

// ChebyshevU returns the Chebyshev polynomial of the second kind
// Uₙ(x) of each element, by the recurrence Uₙ₊₁ = 2x·Uₙ − Uₙ₋₁ with
// U₀ = 1, U₁ = 2x.
func ChebyshevU(n int, x *Array) (*Array, error) {
	if n < 0 {
		return nil, errf("ChebyshevU: degree must be ≥ 0, got %d", n)
	}
	if x.dt == Complex {
		return nil, errf("ChebyshevU: complex arrays are not supported")
	}
	out := &Array{shape: append([]int{}, x.shape...), dt: Float}
	out.alloc(x.Len())
	engine.Parallel(x.Len(), func(s, e int) {
		for i := s; i < e; i++ {
			xv := x.floatAt(i)
			u0, u1 := 1.0, 2*xv
			for k := 1; k < n; k++ {
				u0, u1 = u1, 2*xv*u1-u0
			}
			if n == 0 {
				u1 = 1
			}
			out.floats[i] = u1
		}
	})
	return out, nil
}

// SphericalHarmonicReal returns the real spherical harmonic
// Y_l^m(θ, φ) built from the complex one, both branches under the
// standard convention Y_real = √2·(−1)^m·Re(Y_l^m) for m > 0 and
// Y_real = √2·(−1)^m·Im(Y_l^|m|) for m < 0, and for m = 0 the plain
// Y_l⁰. Values are real and element-wise over theta and phi, which
// must share a shape.
func SphericalHarmonicReal(l, m int, theta, phi *Array) (*Array, error) {
	// m = 0 is purely real in the complex form, but the contract here
	// is a Float array, so the value is extracted rather than handing
	// the complex Y_l⁰ through unchanged.
	yc, err := SphericalHarmonic(l, absInt(m), theta, phi)
	if err != nil {
		return nil, err
	}
	out := &Array{shape: append([]int{}, theta.shape...), dt: Float}
	out.alloc(theta.Len())
	root2 := math.Sqrt2
	// The standard real forms carry the (−1)^m phase on both sides of
	// zero, mirroring the Condon-Shortley factor the complex harmonics
	// carry.
	sign := 1.0
	if m%2 != 0 {
		sign = -1
	}
	for i := range theta.Len() {
		z := yc.complexAt(i)
		switch {
		case m == 0:
			out.floats[i] = real(z)
		case m > 0:
			out.floats[i] = sign * root2 * real(z)
		default:
			out.floats[i] = sign * root2 * imag(z)
		}
	}
	return out, nil
}

// Airy returns the Airy functions of the first and second kind at
// each element, Ai in the first return slot and Bi in the second.
// The power series about zero is used, which converges for every
// real x but is numerically dependable only within |x| ≤ 8; outside
// that window the functions return NaN rather than a silently wrong
// value. Positive arguments above 2 sum the series in extended
// precision, where the float64 reconstruction of Ai would cancel the
// growing solution and lose digits. The scaling convention is the
// standard Ai, Bi with Ai(0) = 1/(3^{2/3}Γ(2/3)),
// Bi(0) = 1/(3^{1/6}Γ(1/3)).
func Airy(a *Array) (ai, bi *Array, err error) {
	if a.dt == Complex {
		return nil, nil, errf("Airy: complex arrays are not supported")
	}
	ai = &Array{shape: append([]int{}, a.shape...), dt: Float}
	ai.alloc(a.Len())
	bi = &Array{shape: append([]int{}, a.shape...), dt: Float}
	bi.alloc(a.Len())
	engine.Parallel(a.Len(), func(s, e int) {
		for i := s; i < e; i++ {
			x := a.floatAt(i)
			if math.Abs(x) > 8 {
				ai.floats[i] = math.NaN()
				bi.floats[i] = math.NaN()
				continue
			}
			aiv, biv := airySeries(x)
			ai.floats[i] = aiv
			bi.floats[i] = biv
		}
	})
	return ai, bi, nil
}

// airySeries evaluates the Maclaurin pair
// Ai(x) = c1·f(x) − c2·g(x), Bi(x) = √3·(c1·f(x) + c2·g(x))
// with f = Σ 3ᵏ(1/3)ₖx^{3k}/(3k)!, g = Σ 3ᵏ(2/3)ₖx^{3k+1}/(3k+1)!,
// summed by the term-ratio recurrence to machine precision. For
// positive arguments above airyBigCut the Ai reconstruction runs in
// extended precision instead: Bi is a sum of like-signed terms and
// keeps the float64 series.
func airySeries(x float64) (aiv, biv float64) {
	const (
		c1 = 0.35502805388781723926
		c2 = 0.25881940379280679840
	)
	tf, tg := 1.0, x
	sumF, sumG := 1.0, x
	for k := range 400 {
		tf *= 3 * x * x * x * (float64(k) + 1.0/3) / float64((3*k+1)*(3*k+2)*(3*k+3))
		tg *= 3 * x * x * x * (float64(k) + 2.0/3) / float64((3*k+2)*(3*k+3)*(3*k+4))
		sumF += tf
		sumG += tg
		if math.Abs(tf) < 1e-18*math.Abs(sumF) && math.Abs(tg) < 1e-18*math.Abs(sumG) {
			break
		}
	}
	biv = math.Sqrt(3) * (c1*sumF + c2*sumG)
	if x > airyBigCut {
		aiv = airyBig(x)
		return aiv, biv
	}
	aiv = c1*sumF - c2*sumG
	return aiv, biv
}

// airyBigCut is the positive argument above which Ai is reconstructed
// in extended precision. At the cut the float64 difference c1·f − c2·g
// has lost 2ζ/ln2 ≈ 5 bits to the cancellation, which caps its
// relative accuracy at about 1e-14; below the cut the float64 series
// still holds fourteen digits.
const airyBigCut = 2.0

// The Maclaurin constants of the extended-precision branch, Ai(0) and
// −Ai'(0) to 50 significant digits: at the working sizes the
// cancellation needs, the float64 roundings of airySeries would
// reintroduce the very error the branch removes. The regression test
// pins their float64 roundings against the float64 constants above.
const (
	airyC1Digits = "0.35502805388781723926006318600418317639797917419918"
	airyC2Digits = "0.25881940379280679840518356018920396347909113835493"
)

// airyBig evaluates Ai(x) for x > 0 in extended precision. The
// reconstruction c1·f − c2·g cancels e^{2ζ}, ζ = 2x^{3/2}/3: both
// products carry the growing solution e^{ζ} while Ai itself is the
// decaying e^{−ζ} one, so the float64 difference loses 2ζ/ln2 ≈
// 1.92·x^{3/2} bits (4.5e-3 relative at x = 8). The same two series are
// therefore summed in math/big with a working size scaled to that loss,
// following the fresnelBig pattern; the term ratios here are the
// simplified forms x³/((3k+2)(3k+3)) and x³/((3k+3)(3k+4)) the float64
// recurrence above evaluates with the (3k+1) and (3k+2) factors still
// in place.
func airyBig(x float64) float64 {
	bits := uint(64 + int(1.9235*x*math.Sqrt(x)) + 96)
	xf := new(big.Float).SetPrec(bits).SetFloat64(x)
	x3 := new(big.Float).SetPrec(bits).Mul(xf, xf)
	x3.Mul(x3, xf)
	c1, _, err := big.ParseFloat(airyC1Digits, 10, bits, big.ToNearestEven)
	if err != nil {
		return math.NaN()
	}
	c2, _, err := big.ParseFloat(airyC2Digits, 10, bits, big.ToNearestEven)
	if err != nil {
		return math.NaN()
	}
	tf := new(big.Float).SetPrec(bits).SetInt64(1)
	tg := new(big.Float).SetPrec(bits).Set(xf)
	sumF := new(big.Float).SetPrec(bits).Set(tf)
	sumG := new(big.Float).SetPrec(bits).Set(tg)
	// The terms are all positive for x > 0, so a cut in absolute size
	// relative to a sum that stays above 1 is a relative cut.
	tiny := new(big.Float).SetPrec(bits).SetMantExp(big.NewFloat(1), -int(bits)+20)
	for k := range 400 {
		r1 := new(big.Float).SetPrec(bits).Quo(x3, big.NewFloat(float64((3*k+2)*(3*k+3))))
		r2 := new(big.Float).SetPrec(bits).Quo(x3, big.NewFloat(float64((3*k+3)*(3*k+4))))
		tf.Mul(tf, r1)
		tg.Mul(tg, r2)
		sumF.Add(sumF, tf)
		sumG.Add(sumG, tg)
		if tf.Cmp(tiny) < 0 && tg.Cmp(tiny) < 0 {
			break
		}
	}
	term := new(big.Float).SetPrec(bits).Mul(c2, sumG)
	aiv := new(big.Float).SetPrec(bits).Mul(c1, sumF)
	aiv.Sub(aiv, term)
	out, _ := aiv.Float64()
	return out
}

// absInt returns |n|.
func absInt(n int) int {
	if n < 0 {
		return -n
	}
	return n
}