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

package stats

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

import (
	"math"
)

// Noncentral distributions: the χ², t and F laws with a noncentrality
// parameter, the laws the power of every test in this package runs on.
// The χ² is the Poisson mixture of its central family, the exact
// identity a noncentral χ²(ν, λ) draw is χ²(ν + 2J) with J a
// Poisson(λ/2) count; the F mixes the same numerator against the
// central denominator, whose pieces carry the scale (ν₁+2i)/ν₁. The
// noncentral t runs Lenth's algorithm
// (Applied Statistics 38, 1989, pages 185 to 189): the CDF as a sum of
// even terms, the Poisson-weighted I_x(j+½, ν/2) of the folded |t|,
// and odd terms, the half-normal-weighted I_x(j+1, ν/2) that carry the
// sign of δ, over x = t²/(t²+ν). Every series here sums positive
// well-scaled terms until the remaining Poisson mass bounds the
// truncation below the floor, and reports an error rather than a
// silently truncated value if the budget runs out. The test file holds
// all three against a direct quadrature of E[Φ(t√(V/ν) − δ)] and
// against the closed forms of the degenerate corners.

// noncentralTermFloor bounds the weight a mixture term may still carry
// when the sum stops: the remaining Poisson mass is below it, so the
// omitted tail cannot reach the 15th digit of the answer.
const noncentralTermFloor = 1e-18

// maxNoncentralTerms bounds the mixture loops. The weights peak at the
// index ⌊λ/2⌋ and the walk needs the peak plus a few standard
// deviations of Poisson spread to cross it, so the budget carries
// noncentralities up to roughly 2·10⁵ in the χ² and F and δ up to
// about 440 in the t; beyond that the refusal is explicit, and so is
// every weight the format cannot hold: they are computed term by term
// in log space, never climbed from a seed that could underflow to
// zero and take the whole sum with it.
const maxNoncentralTerms = 100000

// noncentralPoissonWeight is the Poisson(half) weight of the index i,
// computed term by term in log space. The multiplicative climb from
// the e^{−half} seed the series definitions start from underflows to
// an exact zero once half passes about 745, and a zero seed never
// recovers: every later weight would stay zero while the loop believed
// it had converged. Evaluating each weight from its own logarithm
// keeps the terms near the peak exact at any half the budget can walk
// past, and the genuinely negligible ones answer zero, which is what
// they are.
func noncentralPoissonWeight(half float64, i int) float64 {
	if half == 0 {
		if i == 0 {
			return 1
		}
		return 0
	}
	return math.Exp(-half + float64(i)*math.Log(half) - logGamma(float64(i)+1))
}

// noncentralBudgetRefused reports the explicit refusal when the weight
// peak of a Poisson(half) mixture sits past the term budget, where the
// walk would stop early with a wrong answer instead.
func noncentralBudgetRefused(name string, half float64) error {
	return base.Errf("%s: the weight peak at %d needs a walk past the %d-term budget; the mixture answers only up to that noncentrality",
		name, int(math.Floor(half)), maxNoncentralTerms)
}

// noncentralPeakInsideBudget reports whether the Poisson weight peak
// at ⌊half⌋ plus its spread sits inside the term budget.
func noncentralPeakInsideBudget(half float64) bool {
	peak := math.Floor(half)
	return float64(maxNoncentralTerms) >= peak+8*math.Sqrt(peak)+2
}

// noncentralOddWeight is the j-th half-normal weight of Lenth's odd
// series, δ·λ^j·p_0/(√(2π)·(2j+1)!!) with λ = δ² and p_0 the j = 0
// Poisson weight, carrying the sign of δ. The double factorial comes
// out of its log-space form (2j+1)!! = 2^{j+1}Γ(j+3/2)/√π, so the
// weight is computed from its own logarithm like the even part's and
// underflows only once it is genuinely negligible.
func noncentralOddWeight(shift, half float64, j int) float64 {
	if shift == 0 {
		return 0
	}
	lambda := shift * shift
	if lambda == 0 {
		return 0
	}
	ln := math.Log(math.Abs(shift)) + float64(j)*math.Log(lambda) - half -
		(float64(j)+1.5)*math.Ln2 - logGamma(float64(j)+1.5)
	return math.Copysign(math.Exp(ln), shift)
}

// NoncentralChiSquareCDF returns P(X ≤ x) for X ~ χ²(ν, λ), the
// Poisson(λ/2) mixture of central χ²(ν + 2i) CDFs, each through the
// existing GammaLower. The noncentrality λ must be non-negative and
// finite; λ = 0 answers through ChiSquareCDF exactly.
func NoncentralChiSquareCDF(x float64, df int, lambda float64) (float64, error) {
	const name = "NoncentralChiSquareCDF"
	if df < 1 {
		return 0, base.Errf("%s: df must be ≥ 1, got %d", name, df)
	}
	if math.IsNaN(lambda) || lambda < 0 || math.IsInf(lambda, 0) {
		return 0, base.Errf("%s: lambda must be finite and non-negative, got %g", name, lambda)
	}
	if math.IsNaN(x) {
		return 0, base.Errf("%s: x must be a number, got %g", name, x)
	}
	if x <= 0 {
		return 0, nil
	}
	if lambda == 0 {
		return ChiSquareCDF(x, df)
	}
	return noncentralPoissonMixture(name, df, lambda,
		func(i int) (float64, error) {
			g, err := GammaLower(float64(df)/2+float64(i), x/2)
			if err != nil {
				return 0, base.Errf("%s: %w", name, err)
			}
			return g, nil
		})
}

// NoncentralChiSquareDensity returns the χ²(ν, λ) density at x, the
// same Poisson mixture with central χ² densities, each carrying a
// closed exponential-power form. The support convention gives 0 below
// x = 0; at x = 0 with df = 1 the density is the +Inf the
// noncentralities preserve, with df = 2 it is the finite limit
// e^{−λ/2}/2 of the j = 0 mixture term, and past df = 2 it is 0.
func NoncentralChiSquareDensity(x float64, df int, lambda float64) (float64, error) {
	const name = "NoncentralChiSquareDensity"
	if df < 1 {
		return 0, base.Errf("%s: df must be ≥ 1, got %d", name, df)
	}
	if math.IsNaN(lambda) || lambda < 0 || math.IsInf(lambda, 0) {
		return 0, base.Errf("%s: lambda must be finite and non-negative, got %g", name, lambda)
	}
	if math.IsNaN(x) {
		return 0, base.Errf("%s: x must be a number, got %g", name, x)
	}
	if x < 0 || (x == 0 && df > 2) {
		return 0, nil
	}
	if x == 0 {
		// df = 1: the x^{−½} singularity the mixture integrates; df = 2:
		// the j = 0 term is finite at the origin and its limit
		// e^{−λ/2}/2 is the answer.
		if df == 2 {
			return 0.5 * math.Exp(-lambda/2), nil
		}
		return math.Inf(1), nil
	}
	lambdaHalf := lambda / 2
	if !noncentralPeakInsideBudget(lambdaHalf) {
		return 0, noncentralBudgetRefused(name, lambdaHalf)
	}
	total := 0.0
	for i := range maxNoncentralTerms {
		weight := noncentralPoissonWeight(lambdaHalf, i)
		a := float64(df)/2 + float64(i)
		density := weight * math.Exp(-x/2+(a-1)*math.Log(x)-a*math.Ln2-logGamma(a))
		total += density
		if float64(i) > lambdaHalf+1 && weight < noncentralTermFloor {
			return total, nil
		}
	}
	return 0, base.Errf("%s: the mixture did not converge within %d terms for lambda = %g",
		name, maxNoncentralTerms, lambda)
}

// NoncentralChiSquareQuantile returns the q-quantile of χ²(ν, λ) by
// the same bracketed bisection the central tables use, seeded near the
// mean ν + λ.
func NoncentralChiSquareQuantile(q float64, df int, lambda float64) (float64, error) {
	if df < 1 {
		return 0, base.Errf("NoncentralChiSquareQuantile: df must be ≥ 1, got %d", df)
	}
	if math.IsNaN(lambda) || lambda < 0 || math.IsInf(lambda, 0) {
		return 0, base.Errf("NoncentralChiSquareQuantile: lambda must be finite and non-negative, got %g", lambda)
	}
	return continuousQuantile("NoncentralChiSquareQuantile", q, float64(df)+lambda/2,
		func(x float64) (float64, error) {
			return NoncentralChiSquareCDF(x, df, lambda)
		}, nil)
}

// NoncentralFCDF returns P(X ≤ x) for X ~ F(ν₁, ν₂, λ), the numerator
// χ²(ν₁, λ) carried against the central denominator: the Poisson(λ/2)
// mixture of the scaled central pieces (ν₁+2i)/ν₁·F(ν₁+2i, ν₂), whose
// beta form sums I_{ν₁x/(ν₁x+ν₂)}((ν₁+2i)/2, ν₂/2) over the weights,
// through the existing BetaIncomplete. The λ = 0 corner is the central
// F exactly.
func NoncentralFCDF(x float64, df1, df2 int, lambda float64) (float64, error) {
	const name = "NoncentralFCDF"
	if df1 < 1 || df2 < 1 {
		return 0, base.Errf("%s: df1 and df2 must be ≥ 1, got %d and %d", name, df1, df2)
	}
	if math.IsNaN(lambda) || lambda < 0 || math.IsInf(lambda, 0) {
		return 0, base.Errf("%s: lambda must be finite and non-negative, got %g", name, lambda)
	}
	if math.IsNaN(x) {
		return 0, base.Errf("%s: x must be a number, got %g", name, x)
	}
	if math.IsInf(x, 1) {
		return 1, nil
	}
	if x <= 0 {
		return 0, nil
	}
	// The beta argument saturates at 1 for an x so large the product
	// ν₁x overflows, which is the CDF's own limit there.
	numerator := float64(df1) * x
	arg := 1.0
	if !math.IsInf(numerator, 1) {
		arg = numerator / (numerator + float64(df2))
	}
	return noncentralPoissonMixture(name, df1, lambda,
		func(i int) (float64, error) {
			p, err := BetaIncomplete(arg, (float64(df1)+2*float64(i))/2, float64(df2)/2)
			if err != nil {
				return 0, base.Errf("%s: %w", name, err)
			}
			return p, nil
		})
}

// NoncentralFQuantile returns the q-quantile of F(ν₁, ν₂, λ) by
// bracketed bisection, seeded at 1 in the neighbourhood of the F
// median.
func NoncentralFQuantile(q float64, df1, df2 int, lambda float64) (float64, error) {
	if df1 < 1 || df2 < 1 {
		return 0, base.Errf("NoncentralFQuantile: df1 and df2 must be ≥ 1, got %d and %d", df1, df2)
	}
	if math.IsNaN(lambda) || lambda < 0 || math.IsInf(lambda, 0) {
		return 0, base.Errf("NoncentralFQuantile: lambda must be finite and non-negative, got %g", lambda)
	}
	return continuousQuantile("NoncentralFQuantile", q, 1, func(x float64) (float64, error) {
		return NoncentralFCDF(x, df1, df2, lambda)
	}, nil)
}

// noncentralPoissonMixture sums w_i·term(i) over the Poisson(λ/2)
// weights w_i, the shared engine of the noncentral χ² and F. All terms
// are positive, so the running sum carries no cancellation; the walk
// stops past the weight peak once the weights have sunk below the
// floor. The weights are computed term by term in log space, so no
// noncentrality the budget can reach underflows the walk.
func noncentralPoissonMixture(name string, df int, lambda float64,
	term func(i int) (float64, error)) (float64, error) {
	half := lambda / 2
	if !noncentralPeakInsideBudget(half) {
		return 0, noncentralBudgetRefused(name, half)
	}
	total := 0.0
	for i := range maxNoncentralTerms {
		weight := noncentralPoissonWeight(half, i)
		t, err := term(i)
		if err != nil {
			return 0, err
		}
		total += t * weight
		if float64(i) > half+1 && weight < noncentralTermFloor {
			return total, nil
		}
	}
	return 0, base.Errf("%s: the mixture did not converge within %d terms for lambda = %g",
		name, maxNoncentralTerms, lambda)
}

// NoncentralTCDF returns P(T ≤ t) for T ~ t(ν, δ), through Lenth's
// even and odd series: the even part folds the law through |t|, the
// Poisson(δ²/2)-weighted beta ratios of the |t| event, and the odd
// part carries the sign of δ through the √(2/π)δ(δ²)^j/(2j+1)!!
// half-normal weights. The truncation floor is the remaining Poisson
// mass the error bound 2s(xodd − godd) tracks, the bound the published
// algorithm proves. δ = 0 answers through StudentTCDF exactly, and
// t = 0 through the closed corner Φ(−δ).
func NoncentralTCDF(t float64, df int, delta float64) (float64, error) {
	const name = "NoncentralTCDF"
	if df < 1 {
		return 0, base.Errf("%s: df must be ≥ 1, got %d", name, df)
	}
	if math.IsNaN(t) || math.IsInf(t, 0) {
		return 0, base.Errf("%s: t must be finite, got %g", name, t)
	}
	if math.IsNaN(delta) || math.IsInf(delta, 0) {
		return 0, base.Errf("%s: delta must be finite, got %g", name, delta)
	}
	if delta == 0 {
		return StudentTCDF(t, df)
	}
	// The series is derived on t ≥ 0; the reflection F(t; δ) =
	// 1 − F(−t; −δ), an exact identity of the law, covers the rest.
	flipped := false
	magnitude, shift := t, delta
	if t < 0 {
		flipped = true
		magnitude = -t
		shift = -delta
	}
	x := magnitude * magnitude / (magnitude*magnitude + float64(df))
	if x == 0 {
		// t = 0: the value collapses to Φ(−δ) exactly.
		return NormalCDF(-delta), nil
	}
	lambda := shift * shift
	half := lambda / 2
	if !noncentralPeakInsideBudget(half) {
		return 0, noncentralBudgetRefused(name, half)
	}
	// p_j are the Poisson(half) weights of the even part, q_j the
	// half-normal weights of the odd part, q_j = δ·λ^j·p_0/(√(2π)·(2j+1)!!).
	// Both are computed term by term in log space: the multiplicative
	// climb from the e^{−λ/2} seed underflows to an exact zero once
	// |δ| passes about 39, and a zero seed never recovers, which used
	// to answer a silent 0 for the whole law.
	p := 0.5 * noncentralPoissonWeight(half, 0)
	q := noncentralOddWeight(shift, half, 0)
	remaining := 0.5 - p
	a := 0.5
	b := float64(df) / 2
	rxb := math.Pow(1-x, b)
	// ln B(a, b) at a = ½.
	lnBeta := 0.5*math.Log(math.Pi) + logGamma(b) - logGamma(a+b)
	xodd, err := BetaIncomplete(x, a, b)
	if err != nil {
		return 0, base.Errf("%s: %w", name, err)
	}
	// godd and geven are the beta-integral pieces the recurrences peel
	// off xodd and xeven, the subtraction forms of I_x(a+1, b) and
	// I_x(a, b+1): one beta evaluation seeds the whole walk.
	godd := 2 * rxb * math.Exp(a*math.Log(x)-lnBeta)
	xeven := 1 - rxb
	geven := b * x * rxb
	total := p*xodd + q*xeven
	for en := 1.0; en <= maxNoncentralTerms; en++ {
		a++
		xodd -= godd
		xeven -= geven
		godd *= x * (a + b - 1) / a
		geven *= x * (a + b - 0.5) / (a + 0.5)
		p = 0.5 * noncentralPoissonWeight(half, int(en))
		q = noncentralOddWeight(shift, half, int(en))
		remaining -= p
		total += p*xodd + q*xeven
		if bound := 2 * remaining * (xodd - godd); bound <= noncentralTermFloor {
			total += NormalCDF(-shift)
			if flipped {
				total = 1 - total
			}
			return min(1, max(0, total)), nil
		}
	}
	return 0, base.Errf("%s: the series did not converge within %d terms for delta = %g",
		name, maxNoncentralTerms, delta)
}

// NoncentralTQuantile returns the q-quantile of t(ν, δ) by bracketed
// bisection on the signed axis: the law leans towards δ, so the
// bracket grows from the seed in both directions.
func NoncentralTQuantile(q float64, df int, delta float64) (float64, error) {
	if df < 1 {
		return 0, base.Errf("NoncentralTQuantile: df must be ≥ 1, got %d", df)
	}
	if math.IsNaN(delta) || math.IsInf(delta, 0) {
		return 0, base.Errf("NoncentralTQuantile: delta must be finite, got %g", delta)
	}
	return signedQuantile("NoncentralTQuantile", q, math.Abs(delta)+1,
		func(t float64) (float64, error) {
			return NoncentralTCDF(t, df, delta)
		})
}

// signedQuantile inverts a continuous CDF over the whole real axis,
// the signed twin of continuousQuantile: the bracket starts at ±seed
// and doubles outwards until the CDF straddles q, then halves to
// rounding level under the same unconditional convergence.
func signedQuantile(name string, q float64, seed float64,
	cdf func(float64) (float64, error)) (float64, error) {
	// NaN-rejecting on purpose, as in continuousQuantile.
	if !(q >= 0 && q <= 1) {
		return 0, base.Errf("%s: q must lie in [0, 1], got %g", name, q)
	}
	if q == 0 || q == 1 {
		return 0, base.Errf("%s: q = %g has no finite quantile", name, q)
	}
	lo, hi := -seed, seed
	fLo, err := cdf(lo)
	if err != nil {
		return 0, base.Errf("%s: %w", name, err)
	}
	fHi, err := cdf(hi)
	if err != nil {
		return 0, base.Errf("%s: %w", name, err)
	}
	for fLo > q {
		lo *= 2
		if math.IsInf(lo, 0) {
			return 0, base.Errf("%s: failed to bracket q = %g from below", name, q)
		}
		if fLo, err = cdf(lo); err != nil {
			return 0, base.Errf("%s: %w", name, err)
		}
	}
	for fHi < q {
		hi *= 2
		if math.IsInf(hi, 0) {
			return 0, base.Errf("%s: failed to bracket q = %g from above", name, q)
		}
		if fHi, err = cdf(hi); err != nil {
			return 0, base.Errf("%s: %w", name, err)
		}
	}
	converged := false
	for range 4096 {
		mid := (lo + hi) / 2
		if mid == lo || mid == hi {
			converged = true
			break
		}
		f, err := cdf(mid)
		if err != nil {
			return 0, base.Errf("%s: %w", name, err)
		}
		if f < q {
			lo = mid
		} else {
			hi = mid
		}
	}
	if !converged {
		return 0, base.Errf("%s: the bisection for q = %g did not converge", name, q)
	}
	return (lo + hi) / 2, nil
}