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

package core

import (
	"math"
	"math/cmplx"

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

// Special functions for physics, cosmology and chemistry: the gamma
// and beta families, the error function family, Legendre polynomials,
// spherical harmonics and spherical Bessel functions. Element-wise
// functions run on real arrays and promote ints and float32 per the
// ladder; complex inputs are refused. Values outside a function's
// domain follow the IEEE behaviour of the underlying implementation
// (math.Gamma, for example, returns +Inf at non-positive integers and
// NaN where undefined).

// Gamma returns the gamma function Γ(x) of each element.
func Gamma(a *Array) (*Array, error) { return a.realFunc("Gamma", math.Gamma) }

// LnGamma returns the natural logarithm of |Γ(x)| of each element. The
// sign of Γ for negative arguments is dropped; callers that need it
// evaluate Gamma directly.
func LnGamma(a *Array) (*Array, error) {
	return a.realFunc("LnGamma", func(x float64) float64 {
		lg, _ := math.Lgamma(x)
		return lg
	})
}

// Beta returns the Euler beta function B(x, y) = Γ(x)Γ(y)/Γ(x+y)
// element-wise. Both arrays must have the same shape.
func Beta(x, y *Array) (*Array, error) {
	if !sameShape(x.shape, y.shape) {
		return nil, errf("Beta: shape mismatch %s vs %s", shapeText(x.shape), shapeText(y.shape))
	}
	if x.dt == Complex || y.dt == Complex {
		return nil, errf("Beta: 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, yv := x.floatAt(i), y.floatAt(i)
			// In log space: the Γ product overflows to Inf well before
			// B itself leaves the float64 range (B(100, 80) ≈ 3e-67),
			// and Inf/Inf would turn the answer into NaN.
			lx, sx := math.Lgamma(xv)
			ly, sy := math.Lgamma(yv)
			lxy, sxy := math.Lgamma(xv + yv)
			sign := sx * sy * sxy
			out.floats[i] = float64(sign) * math.Exp(lx+ly-lxy)
		}
	})
	return out, nil
}

// Erf and Erfc return the error function and the complementary error
// function of each element.
func Erf(a *Array) (*Array, error)  { return a.realFunc("Erf", math.Erf) }
func Erfc(a *Array) (*Array, error) { return a.realFunc("Erfc", math.Erfc) }

// Sinc returns the normalised sinc sin(πx)/(πx) of each element, with
// Sinc(0) = 1. The π-normalised convention is the one signal
// processing and interpolation use.
func Sinc(a *Array) (*Array, error) {
	return a.realFunc("Sinc", func(x float64) float64 {
		if x == 0 {
			return 1
		}
		return math.Sin(math.Pi*x) / (math.Pi * x)
	})
}

// Cosm1 returns cos(x) − 1 of each element, kept accurate where the
// direct subtraction loses: as a float64 subtraction cos(x) − 1 has
// no correct significant bit under about |x| ≈ 1e-8, while the factored
// even series in x² holds full precision down to the smallest
// representable argument, where the answer is exactly −x²/2 as far as
// the format can see it. The standard library carries log1p and expm1
// for the logarithmic and exponential side of the same problem and no
// cosinusoidal counterpart, which is the gap this fills.
func Cosm1(a *Array) (*Array, error) {
	return a.realFunc("Cosm1", cosm1)
}

// cosm1Crossover splits the series and the direct subtraction: below
// it the factored series ends well past the rounding floor, above it
// the cancellation the subtraction suffers still leaves every bit the
// answer has.
const cosm1Crossover = math.Pi / 4

// cosm1 evaluates cos(x) − 1 at one point. With t = x² the series
// cos x − 1 = −(t/2)·(1 − t/12·(1 − t/30·(1 − t/56·(1 − t/90·
// (1 − t/132·(1 − t/182·(1 − t/240))))))) truncates around 2e-18
// relative at the crossover and falls quadratically below it; the
// leading −t/2 factor carries the sign, so negative arguments need no
// branch.
func cosm1(x float64) float64 {
	t := x * x
	if math.Abs(x) >= cosm1Crossover {
		return math.Cos(x) - 1
	}
	return -0.5 * t * (1 - t/12*(1-t/30*(1-t/56*(1-t/90*(1-t/132*(1-t/182*(1-t/240)))))))
}

// Legendre returns the Legendre polynomial P_l(x) of each element,
// computed by the Bonnet recurrence. The degree l must be ≥ 0; the
// convention is P_0 = 1, P_1 = x, no extra scaling.
func Legendre(l int, x *Array) (*Array, error) {
	if l < 0 {
		return nil, errf("Legendre: degree must be ≥ 0, got %d", l)
	}
	if x.dt == Complex {
		return nil, errf("Legendre: 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)
			p0, p1 := 1.0, xv
			for k := 1; k < l; k++ {
				p0, p1 = p1, ((2*float64(k)+1)*xv*p1-float64(k)*p0)/float64(k+1)
			}
			if l == 0 {
				p1 = 1
			}
			out.floats[i] = p1
		}
	})
	return out, nil
}

// LegendreAssociated returns the associated Legendre function
// P_l^m(x) of each element, with the Condon-Shortley phase (−1)^m
// folded in and no extra normalisation. Requires |m| ≤ l; a negative m
// uses the standard relation P_l^(−m) = (−1)^m (l−m)!/(l+m)! P_l^m.
func LegendreAssociated(l, m int, x *Array) (*Array, error) {
	if l < 0 {
		return nil, errf("LegendreAssociated: degree must be ≥ 0, got %d", l)
	}
	if m < -l || m > l {
		return nil, errf("LegendreAssociated: order %d out of range for degree %d", m, l)
	}
	if x.dt == Complex {
		return nil, errf("LegendreAssociated: complex arrays are not supported")
	}
	am, neg := m, false
	if am < 0 {
		am, neg = -m, true
	}
	mm := float64(am)
	// Seed sign from the Condon-Shortley phase (−1)^m.
	seedSign := 1.0
	if am%2 == 1 {
		seedSign = -1
	}
	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)
			// P_m^m = (−1)^m (2m−1)!! (1−x²)^(m/2).
			fact := 1.0
			for k := 1; k <= am; k++ {
				fact *= float64(2*k - 1)
			}
			pmm := seedSign * fact * math.Pow(1-xv*xv, mm/2)
			if l == am {
				out.floats[i] = pmm
				continue
			}
			// P_{m+1}^m = x (2m+1) P_m^m, then the fixed-order
			// recurrence up the degree.
			pmp1 := xv * (2*mm + 1) * pmm
			pl := pmp1
			for k := am + 2; k <= l; k++ {
				pl = ((2*float64(k)-1)*xv*pmp1 - float64(k+am-1)*pmm) / float64(k-am)
				pmm, pmp1 = pmp1, pl
			}
			out.floats[i] = pl
		}
	})
	if neg {
		// P_l^(−m) = (−1)^m (l−m)!/(l+m)! P_l^m.
		ratio := math.Exp(LnFactorial(l-am) - LnFactorial(l+am))
		engine.Parallel(out.Len(), func(s, e int) {
			for i := s; i < e; i++ {
				out.floats[i] *= seedSign * ratio
			}
		})
	}
	return out, nil
}

// LnFactorial returns the natural logarithm of n!.
func LnFactorial(n int) float64 {
	if n < 0 {
		return math.NaN()
	}
	lg, _ := math.Lgamma(float64(n + 1))
	return lg
}

// SphericalHarmonic returns the complex spherical harmonic
// Y_l^m(θ, φ) element-wise over the polar angle θ and azimuth φ,
// using the Condon-Shortley phase and the orthonormal convention
// N·P_l^m(cos θ)·e^{imφ} with N = √((2l+1)/(4π)·(l−m)!/(l+m)!).
// theta and phi must have the same shape.
func SphericalHarmonic(l, m int, theta, phi *Array) (*Array, error) {
	if l < 0 {
		return nil, errf("SphericalHarmonic: degree must be ≥ 0, got %d", l)
	}
	if m < -l || m > l {
		return nil, errf("SphericalHarmonic: order %d out of range for degree %d", m, l)
	}
	if !sameShape(theta.shape, phi.shape) {
		return nil, errf("SphericalHarmonic: shape mismatch %s vs %s",
			shapeText(theta.shape), shapeText(phi.shape))
	}
	if theta.dt == Complex || phi.dt == Complex {
		return nil, errf("SphericalHarmonic: complex arrays are not supported")
	}
	am, conj := m, false
	if am < 0 {
		am, conj = -m, true
	}
	// P_l^m takes the cosine of the polar angle, not the angle.
	cosTheta, err := Cos(theta)
	if err != nil {
		return nil, err
	}
	p, err := LegendreAssociated(l, am, cosTheta)
	if err != nil {
		return nil, err
	}
	norm := math.Sqrt((2*float64(l) + 1) / (4 * math.Pi) *
		math.Exp(LnFactorial(l-am)-LnFactorial(l+am)))
	// The Condon-Shortley phase (−1)^m is already folded into the
	// associated Legendre function; the normalisation carries no
	// extra sign.
	out := &Array{shape: append([]int{}, theta.shape...), dt: Complex}
	out.alloc(theta.Len())
	engine.Parallel(theta.Len(), func(s, e int) {
		for i := s; i < e; i++ {
			y := complex(norm*p.floatAt(i), 0) *
				cmplx.Exp(complex(0, float64(am)*phi.floatAt(i)))
			if conj {
				// Y_{l,−m} = (−1)^m conj(Y_{l,m}).
				phase := 1.0
				if am%2 == 1 {
					phase = -1
				}
				y = complex(phase, 0) * cmplx.Conj(y)
			}
			out.complexes[i] = y
		}
	})
	return out, nil
}

// SphericalBesselJ returns the spherical Bessel function of the first
// kind j_l(x) of each element. Orders at or below |x| climb the upward
// recurrence from the exact closed forms j₀ = sin x/x and
// j₁ = sin x/x² − cos x/x, the stable direction there at O(l) cost;
// higher orders run the downward Miller recurrence, which is what the
// small roots need and whose start sits above the turning point at
// order l + |x|.
func SphericalBesselJ(l int, x *Array) (*Array, error) {
	if l < 0 {
		return nil, errf("SphericalBesselJ: degree must be ≥ 0, got %d", l)
	}
	if x.dt == Complex {
		return nil, errf("SphericalBesselJ: 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++ {
			out.floats[i] = sphericalBesselJ(l, x.floatAt(i))
		}
	})
	return out, nil
}

// SphericalBesselY returns the spherical Bessel function of the second
// kind y_l(x) of each element, evaluated by the upward recurrence from
// y_0 = −cos(x)/x, the stable direction. Every y_l diverges to −Inf at
// x = 0.
func SphericalBesselY(l int, x *Array) (*Array, error) {
	if l < 0 {
		return nil, errf("SphericalBesselY: degree must be ≥ 0, got %d", l)
	}
	if x.dt == Complex {
		return nil, errf("SphericalBesselY: 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)
			if xv == 0 {
				// Divergence at the origin: fill this element and carry
				// on, the rest of the chunk still needs its values.
				out.floats[i] = math.Inf(-1)
				continue
			}
			y0 := -math.Cos(xv) / xv
			if l == 0 {
				// Assign and carry on: a return here would abandon the
				// rest of the worker's chunk at zero.
				out.floats[i] = y0
				continue
			}
			y1 := -math.Cos(xv)/(xv*xv) - math.Sin(xv)/xv
			for k := 1; k < l; k++ {
				y0, y1 = y1, (2*float64(k)+1)/xv*y1-y0
			}
			out.floats[i] = y1
		}
	})
	return out, nil
}

// sphericalBesselJ evaluates j_l at one point. Below the argument the
// upward recurrence climbs from the exact j₀ and j₁ at O(l) cost; at
// and above it the downward Miller recurrence, started well above the
// turning point, is the stable direction.
func sphericalBesselJ(l int, x float64) float64 {
	if x == 0 {
		if l == 0 {
			return 1
		}
		return 0
	}
	if math.Abs(x) < 1e-4 {
		// Small-x power series. Below ~1e-6 the unscaled Miller seed
		// overflows on the way down and the renormalisation turns into
		// Inf/Inf = NaN; the two-term series is relative ~1e-8 exact
		// over the whole branch and underflows to the true limit for
		// large l.
		num := math.Pow(x, float64(l))
		dbl := 1.0 // (2l+1)!!
		for k := 3; k <= 2*l+1; k += 2 {
			dbl *= float64(k)
		}
		return num / dbl * (1 - x*x/(2*float64(2*l+3)))
	}
	if float64(l) <= math.Abs(x) {
		return sphericalJUpward(l, x)
	}
	// |x| keeps the start above l for negative arguments too: a start
	// at or below l would never pass the order on the way down and
	// renormalise against a garbage seed.
	ceiling := l + int(math.Abs(x)) + 40
	jp1, j := 0.0, 1.0 // j_{k+1}, j_k, seeded at k = ceiling
	jl := 0.0
	for k := ceiling; k >= 1; k-- {
		if k == l {
			jl = j
		}
		jp1, j = j, (2*float64(k)+1)/x*j-jp1
		if aj := math.Abs(j); aj > 1e200 {
			// The unscaled walk overflows for degree-argument
			// combinations whose true value is representable; the
			// renormalisation cancels any common factor.
			jp1 /= aj
			j /= aj
			jl /= aj
		}
	}
	// j now holds the unscaled j_0, which for l = 0 is the answer.
	if l == 0 {
		jl = j
	}
	return sphericalJ0(x) * jl / j
}

// sphericalJUpward climbs the spherical Bessel functions of the first
// kind from the exact closed forms j₀ = sin x/x and
// j₁ = sin x/x² − cos x/x by the upward recurrence
// jₖ₊₁ = (2k+1)/x·jₖ − jₖ₋₁, the stable direction while the order
// stays at or below the argument. Both seeds carry the absolute scale,
// so no renormalisation is needed, and they carry the parity for free:
// the recurrence reproduces j_l(−x) = (−1)^l j_l(x). The caller
// guarantees x ≠ 0 and 0 ≤ l ≤ |x|.
func sphericalJUpward(l int, x float64) float64 {
	j0 := sphericalJ0(x)
	j1 := math.Sin(x)/(x*x) - math.Cos(x)/x
	if l == 0 {
		return j0
	}
	for k := 1; k < l; k++ {
		j0, j1 = j1, (2*float64(k)+1)/x*j1-j0
	}
	return j1
}

// sphericalJ0 is the exact j_0 = sin(x)/x with the x = 0 limit.
func sphericalJ0(x float64) float64 {
	if x == 0 {
		return 1
	}
	return math.Sin(x) / x
}