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

package core

import "math"

// The exponential-integral and polygamma family. Element-wise over
// real arrays; ints and float32 promote, complex inputs are refused.

// ExpIntegralE1 returns the exponential integral E1(x) =
// ∫ₓ^∞ e^(−t)/t dt of each element, defined for x > 0. A power series
// runs for x ≤ 2 and a continued fraction above, both accurate to
// near machine precision.
func ExpIntegralE1(a *Array) (*Array, error) {
	return a.realFunc("ExpIntegralE1", expIntegralE1)
}

func expIntegralE1(x float64) float64 {
	if x <= 0 {
		return math.NaN()
	}
	if x <= 2 {
		// E1(x) = −γ − ln x − Σ_{k≥1} (−x)^k/(k·k!).
		sum := 0.0
		term := 1.0
		for k := 1; k < 200; k++ {
			term = -term * x / float64(k)
			add := term / float64(k)
			sum += add
			if math.Abs(add) < 1e-18*math.Abs(sum) {
				break
			}
		}
		return -eulerGamma - math.Log(x) - sum
	}
	// Gauss continued fraction of the incomplete gamma function at
	// a = 0, evaluated bottom-up from a generous depth: E1 = e^(-x)/f
	// with f = x+1 - 1^2/(x+3 - 2^2/(x+5 - 3^2/(x+7 - ...))). The
	// numerators grow quadratically, so a hundred levels leave the tail
	// far below double rounding for every x in this branch.
	const cfDepth = 100
	f := x + 2*float64(cfDepth) + 1
	for i := cfDepth; i >= 1; i-- {
		f = (x + 2*float64(i) - 1) - float64(i*i)/f
	}
	return math.Exp(-x) / f
}

// ExpIntegralEi returns the exponential integral Ei(x) of each
// element, defined for real x ≠ 0 by the Cauchy principal value.
// Negative x routes through E1; positive x uses the convergent series
// with the logarithmic singularity subtracted.
func ExpIntegralEi(a *Array) (*Array, error) {
	return a.realFunc("ExpIntegralEi", func(x float64) float64 {
		if x < 0 {
			return -expIntegralE1(-x)
		}
		if x == 0 {
			return math.Inf(-1)
		}
		if x >= 40 {
			// Asymptotic expansion summed to its smallest term at
			// k ≈ x. The convergent series needs k proportional to x
			// (already past 300 terms at x = 250), so the old fixed
			// cap silently truncated it; from x = 40 on, the omitted
			// asymptotic tail sits below double rounding. The exp/log
			// shift evaluates e^x/x without overflowing before the
			// true value does.
			t, sum := 1.0, 1.0
			for k := 1; k <= int(x); k++ {
				t *= float64(k) / x
				sum += t
			}
			return math.Exp(x-math.Log(x)) * sum
		}
		sum := 0.0
		term := 1.0
		for k := 1; k < 300; k++ {
			term *= x / float64(k)
			add := term / float64(k)
			sum += add
			if math.Abs(add) < 1e-18*math.Abs(sum) {
				break
			}
		}
		return eulerGamma + math.Log(x) + sum
	})
}

// Digamma returns the digamma function ψ(x) = d/dx ln Γ(x) of each
// element. A recurrence lifts the argument above 20, then the
// asymptotic series applies; the reflection formula covers x < 0.
func Digamma(a *Array) (*Array, error) {
	return a.realFunc("Digamma", digamma)
}

func digamma(x float64) float64 {
	if math.IsNaN(x) || math.IsInf(x, 0) {
		return math.NaN()
	}
	result := 0.0
	if x <= 0 && x == math.Floor(x) {
		return math.NaN()
	}
	if x < 0 {
		// Reflection: ψ(1−x) = ψ(x) + π·cot(πx).
		result = -math.Pi / math.Tan(math.Pi*x)
		x = 1 - x
	}
	for x < 20 {
		result -= 1 / x
		x++
	}
	inv := 1 / x
	inv2 := inv * inv
	result += math.Log(x) - 0.5*inv - inv2*(1.0/12-inv2*(1.0/120-inv2*(1.0/252-inv2*(1.0/240))))
	return result
}

// Trigamma returns the trigamma function ψ′(x), the derivative of the
// digamma function, of each element. The same recurrence and
// asymptotic strategy as Digamma applies.
func Trigamma(a *Array) (*Array, error) {
	return a.realFunc("Trigamma", trigamma)
}

func trigamma(x float64) float64 {
	if math.IsNaN(x) || math.IsInf(x, 0) {
		return math.NaN()
	}
	result := 0.0
	reflected := false
	if x <= 0 && x == math.Floor(x) {
		return math.NaN()
	}
	if x < 0 {
		// Reflection: ψ′(x) + ψ′(1−x) = π²/sin²(πx), so ψ′(x) is the
		// constant minus the lifted value ψ′(1−x): the recurrence and
		// the asymptotic below subtract instead of add, mirroring the
		// digamma structure above.
		result = math.Pi * math.Pi / (math.Sin(math.Pi*x) * math.Sin(math.Pi*x))
		x = 1 - x
		reflected = true
	}
	for x < 20 {
		if reflected {
			result -= 1 / (x * x)
		} else {
			result += 1 / (x * x)
		}
		x++
	}
	inv := 1 / x
	inv2 := inv * inv
	tail := inv*(1+0.5*inv) + inv*inv2*(1.0/6-inv2*(1.0/30-inv2*(1.0/42-inv2*(1.0/30))))
	if reflected {
		return result - tail
	}
	return result + tail
}

// eulerGamma is the Euler-Mascheroni constant.
const eulerGamma = 0.57721566490153286060651209008240243