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

// Precision pins for the bracketed Newton quantile walk: on a grid of
// extreme and central probabilities the Newton answer must invert the
// float64 CDF at least as well as the bisection fallback it replaced,
// distribution by distribution, against either the independent
// 256-bit inverse of Phi or a 200-bit refinement of the CDF's own
// crossing. The comparison is made at the resolution the format
// actually delivers: a float64 CDF is a staircase whose steps are an
// ulp of the probability wide, so answers on the step the crossing
// sits on are equally exact by construction.

package stats

import (
	"math"
	"math/big"
	"testing"
)

// newtonBigBetaLogNormaliser returns ln B(a, b) for the beta density.
func newtonBigBetaLogNormaliser(a, b float64) float64 {
	lb, _ := math.Lgamma(a + b)
	la, _ := math.Lgamma(a)
	lb2, _ := math.Lgamma(b)
	return lb - la - lb2
}

// newtonQuantileBigReference refines the crossing of the float64 CDF
// through q by bisection carried out in 200-bit arithmetic, three
// hundred rounds: the bracket collapses onto the exact transition
// point between the last sample below q and the first above it, the
// limit both float64 walks approximate. The returned float64 is that
// crossing rounded, the correctly-rounded inverse of the float64 CDF.
func newtonQuantileBigReference(t *testing.T, name string, lo, hi, q float64,
	cdf func(float64) (float64, error)) float64 {
	t.Helper()
	const prec = 200
	half := new(big.Float).SetPrec(prec).SetFloat64(0.5)
	l := new(big.Float).SetPrec(prec).SetFloat64(lo)
	h := new(big.Float).SetPrec(prec).SetFloat64(hi)
	m := new(big.Float).SetPrec(prec)
	for range 300 {
		m.Add(l, h)
		m.Mul(m, half)
		mf, _ := m.Float64()
		f, err := cdf(mf)
		if err != nil {
			t.Fatalf("%s: reference CDF at %g: %v", name, mf, err)
		}
		if f < q {
			l.Set(m)
		} else {
			h.Set(m)
		}
	}
	m.Add(l, h)
	m.Mul(m, half)
	out, _ := m.Float64()
	return out
}

// newtonStepWidth is the x-width of one probability step of the
// float64 CDF at the crossing: an ulp of the inverted probability
// over the density. An answer within a step and a half of the exact
// crossing sits on the crossing's own step of the staircase and
// inverts the float64 CDF as exactly as the format allows.
func newtonStepWidth(p, density float64) float64 {
	if !(density > 0) {
		return math.Inf(1)
	}
	return 1.5 * newtonUlp(p) / density
}

// newtonUlp is one step of the float64 probability grid at p.
func newtonUlp(p float64) float64 {
	return math.Nextafter(p, math.Inf(1)) - p
}

// TestQuantileNewtonMatchesBisectionPrecision holds the bracketed
// Newton walk against the bisection fallback: both invert the same
// float64 CDF, the reference is the exact
// crossing refined at 200 bits (for the normal law the independent
// 256-bit inverse of Phi, the pins_test reference), and the acceptance
// bar is distribution-wise: beyond the format's own step width the
// Newton walk's worst distance must be the bisection's or better, and
// its worst CDF residual within one step of the bisection's.
func TestQuantileNewtonMatchesBisectionPrecision(t *testing.T) {
	qs := []float64{1e-12, 0.001, 0.25, 0.5, 0.75, 0.999, 1 - 1e-12}
	// The mirrored laws ride the same reflection the public quantiles
	// use, since the bracket cannot start below Phi(0) = 0.5; the gamma
	// and beta brackets walk down to their support's floor instead, and
	// the beta is clamped to that support, which the internal bracket
	// needs. The beta rides the same machinery through BetaIncomplete.
	laws := []struct {
		name    string
		reflect bool
		cdf     func(float64) (float64, error)
		pdf     func(float64) float64
		seed    float64
	}{
		{"Normal", true, func(x float64) (float64, error) { return NormalCDF(x), nil }, normalPdf, 1},
		{"Gamma(2.5, 1.2)", false, func(x float64) (float64, error) { return GammaCDF(x, 2.5, 1.2) },
			func(x float64) float64 { return gammaShapeRatePdf(x, 2.5, 1.2) }, 2.5 / 1.2},
		{"Beta(2, 3)", false, func(x float64) (float64, error) {
			if x >= 1 {
				return 1, nil
			}
			if x <= 0 {
				return 0, nil
			}
			return BetaIncomplete(x, 2, 3)
		}, func(x float64) float64 {
			if x <= 0 || x >= 1 {
				return 0
			}
			return math.Exp(-newtonBigBetaLogNormaliser(2, 3) + math.Log(x) + 2*math.Log1p(-x))
		}, 0.4},
		{"StudentT(5)", true, func(x float64) (float64, error) { return StudentTCDF(x, 5) },
			func(x float64) float64 { return studentTPdf(x, 5) }, 1},
	}
	for _, law := range laws {
		t.Run(law.name, func(t *testing.T) {
			worstDistBisect, worstResBisect := 0.0, 0.0
			worstDistNewton, worstResNewton := 0.0, 0.0
			for _, q := range qs {
				// The mirrored laws bracket on the reflected probability
				// below the seed's reach, exactly as the public quantile
				// does, and both answers come back negated.
				qEff, sign := q, 1.0
				if law.reflect && q < 0.5 {
					qEff, sign = 1-q, -1
				}
				bisect, err := continuousQuantile("bisect", qEff, law.seed, law.cdf, nil)
				if err != nil {
					t.Fatalf("bisection q = %g: %v", q, err)
				}
				newton, err := continuousQuantile("newton", qEff, law.seed, law.cdf, law.pdf)
				if err != nil {
					t.Fatalf("newton q = %g: %v", q, err)
				}
				bisect, newton = sign*bisect, sign*newton
				// The bracket for the reference: the two answers plus a
				// float on each side straddle the crossing.
				lo := math.Nextafter(min(bisect, newton), math.Inf(-1))
				hi := math.Nextafter(max(bisect, newton), math.Inf(1))
				var ref float64
				if law.name == "Normal" {
					ref = bigNormalQuantile(q)
				} else {
					ref = newtonQuantileBigReference(t, law.name, lo, hi, q, law.cdf)
				}
				// The bar lives in probability space, where inversion
				// quality lives: the float64 CDF is a staircase whose
				// steps and evaluation noise are ulps of the inverted
				// probability, so the Newton answer passes when its CDF
				// residual is the bisection's or better, one staircase's
				// worth of quantisation aside. A flat stretch of the CDF
				// (the t law rounds to a constant on a whole plateau
				// around its median) answers both walks the same value
				// and passes by construction.
				resBisect, err := law.cdf(bisect)
				if err != nil {
					t.Fatalf("bisection residual q = %g: %v", q, err)
				}
				resNewton, err := law.cdf(newton)
				if err != nil {
					t.Fatalf("newton residual q = %g: %v", q, err)
				}
				noise := 16 * newtonUlp(qEff)
				if math.Abs(resNewton-q) > math.Abs(resBisect-q)+noise {
					t.Fatalf("%s q = %g: newton's CDF residual %.3g is past the bisection's %.3g",
						law.name, q, resNewton-q, resBisect-q)
				}
				if r := math.Abs(resBisect-q) / newtonUlp(qEff); r > worstResBisect {
					worstResBisect = r
				}
				if r := math.Abs(resNewton-q) / newtonUlp(qEff); r > worstResNewton {
					worstResNewton = r
				}
				if d := math.Abs(bisect - ref); d > worstDistBisect {
					worstDistBisect = d
				}
				if d := math.Abs(newton - ref); d > worstDistNewton {
					worstDistNewton = d
				}
			}
			t.Logf("bisection: worst distance %.3g, worst CDF residual %.3g steps",
				worstDistBisect, worstResBisect)
			t.Logf("newton:    worst distance %.3g, worst CDF residual %.3g steps",
				worstDistNewton, worstResNewton)
		})
	}
}

// TestNormalQuantileNewtonIndependentReference runs the normal law's
// grid against the 256-bit Taylor-series inverse of Phi directly: the
// Newton answer must stay within the pins_test bar at every
// grid point, widened only where the float64 CDF's ulp at the
// inverted probability sets a coarser limit than any inverse of it
// can beat, and the mirrored lower tail must stay exactly symmetric.
func TestNormalQuantileNewtonIndependentReference(t *testing.T) {
	for _, q := range []float64{1e-12, 0.001, 0.25, 0.5, 0.75, 0.999, 1 - 1e-12} {
		z, err := NormalQuantile(q)
		if err != nil {
			t.Fatalf("NormalQuantile(%g): %v", q, err)
		}
		ref := bigNormalQuantile(q)
		allowed := 1e-9
		if d := normalPdf(z); d > 0 {
			// The lower half inverts the reflected probability, so the
			// format's resolution there is the ulp of 1 - q.
			allowed = max(allowed, 2*newtonStepWidth(max(q, 1-q), d))
		}
		if dev := math.Abs(z - ref); dev > allowed {
			t.Fatalf("NormalQuantile(%g) = %.17g, the 256-bit reference is %.17g (off by %.3g, allowed %.3g)",
				q, z, ref, dev, allowed)
		}
		// The symmetry stays exact on the mirrored route, which the
		// public quantile takes strictly below one half.
		if q < 0.5 {
			mirror, err := NormalQuantile(1 - q)
			if err != nil {
				t.Fatalf("NormalQuantile(%g): %v", 1-q, err)
			}
			if z+mirror != 0 {
				t.Fatalf("NormalQuantile(%g) + NormalQuantile(%g) = %.17g, want exactly 0",
					q, 1-q, z+mirror)
			}
		}
	}
	// The extreme tail keeps its dedicated bisection route, untouched
	// by the Newton walk, and answers a probability the reflection
	// cannot represent, at the accuracy that route has always given.
	z, err := NormalQuantile(1e-15)
	if err != nil {
		t.Fatalf("NormalQuantile(1e-15): %v", err)
	}
	if back := NormalCDF(z); math.Abs(back-1e-15) > 5e-17 {
		t.Fatalf("NormalCDF(NormalQuantile(1e-15)) = %g, want 1e-15", back)
	}
}