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

package stats

import (
	"math"
	"math/big"
	"testing"
	"time"

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

// The contingency tests against exact referents. The hypergeometric
// and the central binomial probabilities are rational numbers, so the
// Fisher and McNemar p-values are compared against a big.Rat
// evaluation of the same definition, not against a tolerance quote.

// exactFactorial returns n! as a big.Int.
func exactFactorial(n int64) *big.Int {
	out := big.NewInt(1)
	for i := int64(2); i <= n; i++ {
		out.Mul(out, big.NewInt(i))
	}
	return out
}

// exactChoose returns C(n, k) as a big.Int.
func exactChoose(n, k int64) *big.Int {
	return new(big.Int).Div(exactFactorial(n), new(big.Int).Mul(exactFactorial(k), exactFactorial(n-k)))
}

// exactFisherTwoSided evaluates the two-sided Fisher p-value of the
// 2×2 table (a b / c d) exactly: the sum of every table probability
// that does not exceed the observed one's, over the hypergeometric
// support the margins induce.
func exactFisherTwoSided(a, b, c, d int64) *big.Rat {
	row1, row2 := a+b, c+d
	col1 := a + c
	col2 := b + d
	lo := max(int64(0), row1-col2)
	hi := min(row1, col1)
	den := new(big.Rat).SetInt(exactChoose(row1+row2, row1))
	observed := new(big.Rat)
	terms := make([]*big.Rat, 0, hi-lo+1)
	for x := lo; x <= hi; x++ {
		p := new(big.Rat).SetFrac(
			new(big.Int).Mul(exactChoose(col1, x), exactChoose(col2, row1-x)),
			den.Num())
		terms = append(terms, p)
		if x == a {
			observed.Set(p)
		}
	}
	sum := new(big.Rat)
	for _, p := range terms {
		if p.Cmp(observed) <= 0 {
			sum.Add(sum, p)
		}
	}
	return sum
}

// exactMcNemar evaluates the exact McNemar p-value of the discordant
// pair (b, c) exactly: twice the binomial(m, ½) tail at min(b, c).
func exactMcNemar(b, c int64) *big.Rat {
	m := b + c
	if m == 0 {
		return big.NewRat(1, 1)
	}
	k := min(b, c)
	den := new(big.Int).Lsh(big.NewInt(1), uint(m))
	tail := new(big.Rat)
	for i := int64(0); i <= k; i++ {
		term := new(big.Rat).SetFrac(exactChoose(m, i), den)
		tail.Add(tail, term)
	}
	// The doubled tail is clamped at 1 exactly as the entry point
	// clamps it: at the middle pair min(b, c) = m/2 the sum already
	// covers the median and doubling overruns.
	if tail.Cmp(big.NewRat(1, 2)) > 0 {
		return big.NewRat(1, 1)
	}
	return tail.Add(tail, tail)
}

// table2x2 builds a 2×2 float array from four counts.
func table2x2(t *testing.T, a, b, c, d float64) *core.Array {
	t.Helper()
	tab, err := core.FromFloats([]float64{a, b, c, d}, 2, 2)
	if err != nil {
		t.Fatal(err)
	}
	return tab
}

func TestFisherExactAgainstRationalReferent(t *testing.T) {
	tables := [][4]int64{
		{12, 5, 7, 10},
		{8, 2, 1, 5},
		{3, 1, 1, 3},
		{1, 9, 9, 1},
		{0, 6, 6, 0},
		{17, 0, 0, 23},
		{35, 12, 41, 9},
	}
	for _, tab := range tables {
		p, or, err := FisherExactTest(table2x2(t, float64(tab[0]), float64(tab[1]), float64(tab[2]), float64(tab[3])), TwoSided)
		if err != nil {
			t.Fatalf("FisherExactTest(%v): %v", tab, err)
		}
		want := exactFisherTwoSided(tab[0], tab[1], tab[2], tab[3])
		wantF, _ := want.Float64()
		if math.Abs(p-wantF) > 1e-12 {
			t.Fatalf("FisherExactTest(%v) two-sided = %.16f, want the exact %.16f", tab, p, wantF)
		}
		wantOR := float64(tab[0]) * float64(tab[3]) / (float64(tab[1]) * float64(tab[2]))
		if math.IsNaN(wantOR) {
			if !math.IsNaN(or) {
				t.Fatalf("FisherExactTest(%v): odds ratio %g, want NaN", tab, or)
			}
		} else if or != wantOR {
			t.Fatalf("FisherExactTest(%v): odds ratio %g, want %g", tab, or, wantOR)
		}
	}
}

func TestFisherExactTeaTablePin(t *testing.T) {
	// The classical 3/1 versus 1/3 table: the two-sided answer is
	// 34/70 by hand, every table but the central one being at most as
	// probable as the observed one.
	p, or, err := FisherExactTest(table2x2(t, 3, 1, 1, 3), TwoSided)
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(p-34.0/70.0) > 1e-12 {
		t.Fatalf("the tea table p = %.16f, want 34/70 = %.16f", p, 34.0/70.0)
	}
	if or != 9 {
		t.Fatalf("the tea table odds ratio = %g, want 9", or)
	}
}

func TestFisherExactOneSidedTails(t *testing.T) {
	tab := table2x2(t, 12, 5, 7, 10)
	less, _, err := FisherExactTest(tab, Less)
	if err != nil {
		t.Fatal(err)
	}
	greater, _, err := FisherExactTest(tab, Greater)
	if err != nil {
		t.Fatal(err)
	}
	// The observed first cell (12) sits above its expectation 9.5: the
	// greater tail must be the shorter one. The two-sided sum carries
	// no ordering guarantee against the tails, because it collects
	// only the tables at most as probable as the observed one while a
	// one-sided tail sweeps the whole arm including the mode.
	if !(greater < less) {
		t.Fatalf("the greater tail %g does not fall below the less tail %g on an upper-tail table", greater, less)
	}
	// The greater tail, recomputed from the exact binomial
	// coefficients over the support's upper arm: the margins are
	// row1 = 17, row2 = 17, col1 = 19, col2 = 15, the support runs
	// 2..17, and the tail collects x = 12..17.
	sum := new(big.Rat)
	den := new(big.Rat).SetInt(exactChoose(34, 17))
	for x := int64(12); x <= 17; x++ {
		term := new(big.Rat).SetInt(new(big.Int).Mul(exactChoose(19, x), exactChoose(15, 17-x)))
		term.Quo(term, den)
		sum.Add(sum, term)
	}
	want, _ := sum.Float64()
	if math.Abs(greater-want) > 1e-12 {
		t.Fatalf("the greater tail = %.16f, want the exact %.16f", greater, want)
	}
}

func TestFisherExactDegenerateMargins(t *testing.T) {
	// A zero margin admits one table only: p = 1 with the odds ratio
	// the formula speaks.
	p, or, err := FisherExactTest(table2x2(t, 0, 5, 0, 7), TwoSided)
	if err != nil {
		t.Fatal(err)
	}
	if p != 1 {
		t.Fatalf("a zero-row table answered p = %g, want 1", p)
	}
	if !math.IsNaN(or) {
		t.Fatalf("a zero-row table answered odds ratio %g, want NaN", or)
	}
	// A zero cell against a full one inside living margins: the odds
	// ratio goes infinite and the two-sided p-value is the observed
	// corner's own probability, 1/210 by the exact coefficients.
	p, or, err = FisherExactTest(table2x2(t, 6, 0, 0, 4), TwoSided)
	if err != nil {
		t.Fatal(err)
	}
	if math.Abs(p-1.0/210.0) > 1e-12 {
		t.Fatalf("the corner table answered p = %.16f, want 1/210 = %.16f", p, 1.0/210.0)
	}
	if or != math.Inf(1) {
		t.Fatalf("the corner table answered odds ratio %g, want +Inf", or)
	}
}

func TestFisherExactRefusals(t *testing.T) {
	if _, _, err := FisherExactTest(mustStatArray(t, []float64{1, 2, 3}, 3), TwoSided); err == nil {
		t.Fatal("a rank-1 table was accepted")
	}
	if _, _, err := FisherExactTest(table2x2(t, 1, 2, 3, 4), Alternative(7)); err == nil {
		t.Fatal("an unknown alternative was accepted")
	}
	if _, _, err := FisherExactTest(table2x2(t, 1, 2, 3, -4), TwoSided); err == nil {
		t.Fatal("a negative count was accepted")
	}
	if _, _, err := FisherExactTest(table2x2(t, 1, 2.5, 3, 4), TwoSided); err == nil {
		t.Fatal("a fractional count was accepted")
	}
	if _, _, err := FisherExactTest(table2x2(t, 1, math.NaN(), 3, 4), TwoSided); err == nil {
		t.Fatal("a NaN count was accepted")
	}
	// Margins wide enough to push the support past the cap: the refusal
	// names the cap instead of enumerating millions of tables.
	if _, _, err := FisherExactTest(table2x2(t, 3e6, 3e6, 3e6, 3e6), TwoSided); err == nil {
		t.Fatal("an over-cap support was accepted")
	}
}

// mustStatArray builds a float array for the refusal probes.
func mustStatArray(t *testing.T, vals []float64, shape ...int) *core.Array {
	t.Helper()
	a, err := core.FromFloats(vals, shape...)
	if err != nil {
		t.Fatal(err)
	}
	return a
}

func TestMcNemarAgainstRationalReferent(t *testing.T) {
	pairs := [][2]int64{
		{5, 15},
		{0, 10},
		{10, 0},
		{7, 7},
		{1, 1},
		{0, 0},
		{23, 41},
	}
	for _, pr := range pairs {
		p, err := McNemarTest(table2x2(t, 10, float64(pr[0]), float64(pr[1]), 12))
		if err != nil {
			t.Fatalf("McNemarTest(%v): %v", pr, err)
		}
		wantF, _ := exactMcNemar(pr[0], pr[1]).Float64()
		if math.Abs(p-wantF) > 1e-12 {
			t.Fatalf("McNemarTest(b=%d, c=%d) = %.16f, want the exact %.16f", pr[0], pr[1], p, wantF)
		}
	}
	// Large discordant counts: the incomplete-beta route must keep
	// answering well past where an m-term sum would have become too
	// expensive to run.
	b, c := 2000.0, 3000.0
	p, err := McNemarTest(table2x2(t, 100, b, c, 100))
	if err != nil {
		t.Fatal(err)
	}
	if !(0 < p && p < 1) {
		t.Fatalf("McNemarTest(2000, 3000) = %g, want a probability in (0, 1)", p)
	}
	pSmall, err := McNemarTest(table2x2(t, 100, 0, 900, 100))
	if err != nil {
		t.Fatal(err)
	}
	if !(pSmall < 1e-15) {
		t.Fatalf("McNemarTest(0, 900) = %g, want a vanishing tail", pSmall)
	}
}

func TestMcNemarRefusals(t *testing.T) {
	if _, err := McNemarTest(mustStatArray(t, []float64{1, 2, 3, 4, 5, 6}, 2, 3)); err == nil {
		t.Fatal("a 2×3 table was accepted")
	}
	if _, err := McNemarTest(table2x2(t, 1, 2, 3, math.Inf(1))); err == nil {
		t.Fatal("an infinite count was accepted")
	}
}

func TestChiSquareIndependenceTwoByTwoIdentity(t *testing.T) {
	// For a 2×2 table the definitional sum collapses algebraically to
	// n(ad−bc)² over the product of the four margins, an independent
	// route the sum must reproduce bit for bit.
	tables := [][4]float64{
		{12, 7, 5, 16},
		{10, 20, 20, 30},
		{1, 9, 9, 1},
		{153, 71, 44, 232},
	}
	for _, tab := range tables {
		chi2, df, p, err := ChiSquareIndependence(table2x2(t, tab[0], tab[1], tab[2], tab[3]))
		if err != nil {
			t.Fatalf("ChiSquareIndependence(%v): %v", tab, err)
		}
		if df != 1 {
			t.Fatalf("ChiSquareIndependence(%v): df = %d, want 1", tab, df)
		}
		n := tab[0] + tab[1] + tab[2] + tab[3]
		want := n * math.Pow(tab[0]*tab[3]-tab[1]*tab[2], 2) /
			((tab[0] + tab[1]) * (tab[2] + tab[3]) * (tab[0] + tab[2]) * (tab[1] + tab[3]))
		if math.Abs(chi2-want) > 1e-12*math.Max(1, want) {
			t.Fatalf("ChiSquareIndependence(%v): χ² = %.16f, want the closed form %.16f", tab, chi2, want)
		}
		if !(0 < p && p < 1) {
			t.Fatalf("ChiSquareIndependence(%v): p = %g is not a probability in (0, 1)", tab, p)
		}
		// Cramér's V for a 2×2 table is |φ| = |ad−bc| over the square
		// root of the same product of margins.
		v, err := CramersV(table2x2(t, tab[0], tab[1], tab[2], tab[3]))
		if err != nil {
			t.Fatal(err)
		}
		phi := math.Abs(tab[0]*tab[3]-tab[1]*tab[2]) /
			math.Sqrt((tab[0]+tab[1])*(tab[2]+tab[3])*(tab[0]+tab[2])*(tab[1]+tab[3]))
		if math.Abs(v-phi) > 1e-12 {
			t.Fatalf("CramersV(%v) = %.16f, want |φ| = %.16f", tab, v, phi)
		}
	}
}

func TestChiSquareIndependenceLargerTable(t *testing.T) {
	// A 3×3 table with the statistic recomputed from the definition
	// through an independent walk over the expected counts.
	counts := []float64{
		21, 8, 3,
		9, 15, 11,
		4, 12, 30,
	}
	tab := mustStatArray(t, counts, 3, 3)
	chi2, df, p, err := ChiSquareIndependence(tab)
	if err != nil {
		t.Fatal(err)
	}
	if df != 4 {
		t.Fatalf("df = %d, want 4", df)
	}
	rows := make([]float64, 3)
	cols := make([]float64, 3)
	total := 0.0
	for i := range 3 {
		for j := range 3 {
			rows[i] += counts[i*3+j]
			cols[j] += counts[i*3+j]
			total += counts[i*3+j]
		}
	}
	want := 0.0
	for i := range 3 {
		for j := range 3 {
			e := rows[i] * cols[j] / total
			want += math.Pow(counts[i*3+j]-e, 2) / e
		}
	}
	if math.Abs(chi2-want) > 1e-12*math.Max(1, want) {
		t.Fatalf("χ² = %.16f, want %.16f", chi2, want)
	}
	v, err := CramersV(tab)
	if err != nil {
		t.Fatal(err)
	}
	wantV := math.Sqrt(want / (total * 2))
	if math.Abs(v-wantV) > 1e-12 {
		t.Fatalf("Cramér's V = %.16f, want %.16f", v, wantV)
	}
	// Perfect independence answers χ² = 0, p = 1 and V = 0.
	independent := []float64{10, 20, 10, 20, 40, 20, 10, 20, 10}
	chi2, _, p, err = ChiSquareIndependence(mustStatArray(t, independent, 3, 3))
	if err != nil {
		t.Fatal(err)
	}
	if chi2 != 0 || p != 1 {
		t.Fatalf("a perfectly independent table answered χ² = %g, p = %g", chi2, p)
	}
}

func TestChiSquareIndependenceRefusals(t *testing.T) {
	if _, _, _, err := ChiSquareIndependence(mustStatArray(t, []float64{1, 2, 3, 4}, 4)); err == nil {
		t.Fatal("a rank-1 table was accepted")
	}
	if _, _, _, err := ChiSquareIndependence(mustStatArray(t, []float64{1, 2, 3}, 3, 1)); err == nil {
		t.Fatal("a single-column table was accepted")
	}
	if _, _, _, err := ChiSquareIndependence(table2x2(t, 1, 0, 3, 0)); err == nil {
		t.Fatal("a zero-column margin was accepted")
	}
	if _, _, _, err := ChiSquareIndependence(table2x2(t, 1, 2, 3, 2.5)); err == nil {
		t.Fatal("a fractional count was accepted")
	}
	if _, err := CramersV(table2x2(t, 1, 2, 3, -4)); err == nil {
		t.Fatal("a negative count was accepted")
	}
	if _, _, _, err := ChiSquareIndependence(table2x2(t, 0, 0, 0, 0)); err == nil {
		t.Fatal("an all-zero table was accepted")
	}
}

func TestContingencyDtypeAgreement(t *testing.T) {
	// A table carried in an integer dtype must answer the same numbers
	// as the float array of the widened values.
	vals := []float64{12, 5, 7, 10}
	fv, _, err := FisherExactTest(table2x2(t, vals[0], vals[1], vals[2], vals[3]), TwoSided)
	if err != nil {
		t.Fatal(err)
	}
	iv := make([]int64, len(vals))
	for i, v := range vals {
		iv[i] = int64(v)
	}
	ia, err := core.FromInts(iv, 2, 2)
	if err != nil {
		t.Fatal(err)
	}
	pv, or, err := FisherExactTest(ia, TwoSided)
	if err != nil {
		t.Fatal(err)
	}
	if pv != fv || or != 120.0/35.0 {
		t.Fatalf("the int table answered (%g, %g), want (%g, %g)", pv, or, fv, 120.0/35.0)
	}
}

func TestFisherExactLopsidedOverflow(t *testing.T) {
	// A corner table with living margins: the observed table sits
	// hundreds of nats below the hypergeometric mode, so the ratio of
	// the mode's probability to the observed one's overflows the
	// format. The one-sided tail over the whole support is the certain
	// event and must answer exactly 1, the opposite tail and the
	// two-sided sum underflow to the correctly rounded zero, and none
	// of the three may turn into the NaN a quotient of infinities
	// produces.
	greater, _, err := FisherExactTest(table2x2(t, 0, 1000, 1000, 0), Greater)
	if err != nil {
		t.Fatal(err)
	}
	if greater != 1 {
		t.Fatalf("the certain Greater tail answered %g, want 1", greater)
	}
	less, _, err := FisherExactTest(table2x2(t, 0, 1000, 1000, 0), Less)
	if err != nil {
		t.Fatal(err)
	}
	twoSided, _, err := FisherExactTest(table2x2(t, 0, 1000, 1000, 0), TwoSided)
	if err != nil {
		t.Fatal(err)
	}
	if math.IsNaN(less) || math.IsNaN(twoSided) {
		t.Fatalf("the far tails answered NaN (%g, %g)", less, twoSided)
	}
	if less != 0 || twoSided != 0 {
		// 1/C(2000,1000) and 2/C(2000,1000) sit near 1e-600: the
		// correctly rounded float64 of both is 0.
		t.Fatalf("the underflowed tails answered (%g, %g), want (0, 0)", less, twoSided)
	}
	// A table whose two-sided answer lands in the representable
	// subnormal window while the mode's ratio still overflows: the
	// exact referent is 2/C(1040,520), near 2e-313.
	p, _, err := FisherExactTest(table2x2(t, 0, 520, 520, 0), TwoSided)
	if err != nil {
		t.Fatal(err)
	}
	want, _ := new(big.Rat).SetFrac(
		big.NewInt(2),
		exactChoose(1040, 520)).Float64()
	if !(p > 0) || math.Abs(p-want) > 1e-318 {
		t.Fatalf("the subnormal two-sided tail answered %.16e, want the exact %.16e", p, want)
	}
}

func TestFisherExactUnknownAlternativeBeforeDegenerateMargins(t *testing.T) {
	// The alternative is refused whatever the margins: a zero-margin
	// table must not short-circuit the refusal behind a p-value of 1.
	if _, _, err := FisherExactTest(table2x2(t, 0, 5, 0, 7), Alternative(7)); err == nil {
		t.Fatal("an unknown alternative was accepted on a zero-margin table")
	}
}

func TestFisherExactCountBeyondTheFormatsIntegers(t *testing.T) {
	// A count past 2^51 can push the margins past 2^52, where float64
	// holds only every second integer: the support walk's increment
	// stops advancing and the enumeration cannot leave its first table.
	// The table below answers a support of one rounded point, so the
	// cap check cannot catch it; the refusal must, and it must come
	// back rather than hang.
	var err error
	pinWatchdog(t, 10*time.Second, "FisherExactTest on [[1e18,0],[0,3]]", func() {
		_, _, err = FisherExactTest(table2x2(t, 1e18, 0, 0, 3), TwoSided)
	})
	if err == nil {
		t.Fatal("a count past the format's integers was accepted")
	}
	if _, err := McNemarTest(table2x2(t, 1, 1e18, 1, 1)); err == nil {
		t.Fatal("McNemarTest accepted a count past the format's integers")
	}
	// The last count the format holds stepwise stays answerable: the
	// bound refuses nothing a legal enumeration can run on.
	p, _, err := FisherExactTest(table2x2(t, float64(maxExactCount-2), 0, 0, 3), TwoSided)
	if err != nil {
		t.Fatalf("a count at the bound's inside was refused: %v", err)
	}
	if !(p >= 0 && p <= 1) {
		t.Fatalf("the boundary table answered p = %g", p)
	}
}