// 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"
	"sourcedock.dev/petrbalvin/tensor/internal/core"
)

import (
	"math"
	"slices"
)

// Statistical inference on samples: association matrices, group
// comparison, goodness of fit, distribution comparison and the
// percentile bootstrap. Everything here is classical frequentist
// tooling built on the distribution functions of cdf.go, so a
// p-value travels no further than the library's own incomplete
// gamma and beta.

// CovarianceMatrix returns the sample covariance matrix of the
// observations in a, an (n, p) array whose rows are observations and
// columns variables: the (p, p) result has cov(i, j) the sample
// covariance of columns i and j with the 1/(n−1) normalisation. At
// least two observations are needed.
func CovarianceMatrix(a *core.Array) (*core.Array, error) {
	return scatterMatrix("CovarianceMatrix", a, false)
}

// CorrelationMatrix returns the Pearson correlation matrix of the
// observations in a, the covariance matrix normalised by each
// column's sample standard deviation; the diagonal is exactly 1.
func CorrelationMatrix(a *core.Array) (*core.Array, error) {
	return scatterMatrix("CorrelationMatrix", a, true)
}

// scatterMatrix builds the covariance or correlation matrix of an
// (n, p) observation array.
func scatterMatrix(name string, a *core.Array, correlate bool) (*core.Array, error) {
	if a.NDim() != 2 {
		return nil, base.Errf("%s: needs a 2-D array of observations, got shape %s",
			name, base.ShapeText(a.Shape()))
	}
	n, p := a.Shape()[0], a.Shape()[1]
	if n < 2 {
		return nil, base.Errf("%s: at least two observations are needed, got %d", name, n)
	}
	if a.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex observations are not supported", name)
	}
	if err := checkFinite(name, "the observations", a); err != nil {
		return nil, err
	}
	// Column means, then centred column copies. A dense float64 payload
	// is read straight from the raw slice, one dispatch-free pass.
	rows := rawFloats(a)
	means := make([]float64, p)
	for j := range p {
		s := 0.0
		if rows != nil {
			for i := range n {
				s += rows[i*p+j]
			}
		} else {
			for i := range n {
				s += a.FloatAt(i*p + j)
			}
		}
		means[j] = s / float64(n)
	}
	cols := make([][]float64, p)
	for j := range p {
		cols[j] = make([]float64, n)
		if rows != nil {
			for i := range n {
				cols[j][i] = rows[i*p+j] - means[j]
			}
		} else {
			for i := range n {
				cols[j][i] = a.FloatAt(i*p+j) - means[j]
			}
		}
	}
	scale := float64(n - 1)
	cov := make([]float64, p*p)
	stds := make([]float64, p)
	for i := range p {
		for j := i; j < p; j++ {
			s := 0.0
			for k := range n {
				s += cols[i][k] * cols[j][k]
			}
			cov[i*p+j] = s / scale
			cov[j*p+i] = cov[i*p+j]
		}
		if correlate {
			stds[i] = math.Sqrt(cov[i*p+i])
			if stds[i] == 0 {
				return nil, base.Errf("CorrelationMatrix: column %d has zero variance", i)
			}
		}
	}
	if correlate {
		for i := range p {
			for j := i; j < p; j++ {
				c := cov[i*p+j] / (stds[i] * stds[j])
				cov[i*p+j] = c
				cov[j*p+i] = c
			}
			cov[i*p+i] = 1
		}
	}
	return floatsToArray(cov, []int{p, p}), nil
}

// WelchTTest compares two independent samples with Welch's t-test,
// which asks no equal-variance assumption: the statistic is
// (mean(a) − mean(b)) over the pooled standard error of the two
// means, and the degrees of freedom follow the Welch-Satterthwaite
// estimate. Returns the statistic, the (generally fractional) degrees
// of freedom and the two-sided p-value. Both samples need at least
// two observations.
func WelchTTest(a, b *core.Array) (t, df, pValue float64, err error) {
	m1, v1, n1, err := sampleMeanVar(a, "WelchTTest")
	if err != nil {
		return 0, 0, 0, err
	}
	m2, v2, n2, err := sampleMeanVar(b, "WelchTTest")
	if err != nil {
		return 0, 0, 0, err
	}
	se1 := v1 / float64(n1)
	se2 := v2 / float64(n2)
	se := se1 + se2
	if se == 0 {
		return 0, 0, 0, base.Errf("WelchTTest: both samples have zero variance")
	}
	t = (m1 - m2) / math.Sqrt(se)
	df = se * se / (se1*se1/float64(n1-1) + se2*se2/float64(n2-1))
	// Two-sided p = P(T > |t|) + P(T < −|t|) = I_z(df/2, 1/2) with
	// z = df/(df + t²), the Student-t tail in closed form.
	z := df / (df + t*t)
	pValue, err = BetaIncomplete(z, df/2, 0.5)
	if err != nil {
		return 0, 0, 0, base.Errf("WelchTTest: %w", err)
	}
	return t, df, pValue, nil
}

// sampleMeanVar returns the mean and the 1/(n−1) variance of a sample.
func sampleMeanVar(a *core.Array, name string) (mean, variance float64, n int, err error) {
	if a.Dtype() == core.Complex {
		return 0, 0, 0, base.Errf("%s: complex samples are not supported", name)
	}
	n = a.Len()
	if n < 2 {
		return 0, 0, 0, base.Errf("%s: at least two observations are needed, got %d", name, n)
	}
	// The finiteness refusal belongs here, before the arithmetic: a NaN
	// sample would otherwise surface only when the tail function
	// rejects the NaN statistic it produced, under the tail function's
	// own name.
	if err := checkFinite(name, "the sample", a); err != nil {
		return 0, 0, 0, err
	}
	mean = 0.0
	// The walk is bounded by the sample's own element count: a rebased
	// view's payload may run past its visible elements, and those
	// invisible tail slots are nobody's observations.
	if fs := rawFloats(a); fs != nil {
		fs = fs[:n]
		for _, v := range fs {
			mean += v
		}
		mean /= float64(n)
		variance = sqDeviations(fs, mean)
	} else {
		for i := range n {
			mean += a.FloatAt(i)
		}
		mean /= float64(n)
		variance = sqDeviationsAt(a, mean)
	}
	variance /= float64(n - 1)
	return mean, variance, n, nil
}

// ChiSquareGoodnessOfFit runs Pearson's test of an observed frequency
// table against expected ones: the statistic is Σ(O−E)²/E over the
// bins, the degrees of freedom the bin count minus one, and the
// p-value the upper tail of that χ² distribution. Every expected
// entry must be positive.
func ChiSquareGoodnessOfFit(observed, expected *core.Array) (chi2 float64, df int, pValue float64, err error) {
	if observed.Dtype() == core.Complex || expected.Dtype() == core.Complex {
		return 0, 0, 0, base.Errf("ChiSquareGoodnessOfFit: complex arrays are not supported")
	}
	n := observed.Len()
	if expected.Len() != n {
		return 0, 0, 0, base.Errf("ChiSquareGoodnessOfFit: observed has %d bins, expected %d",
			n, expected.Len())
	}
	if n < 2 {
		return 0, 0, 0, base.Errf("ChiSquareGoodnessOfFit: at least two bins are needed, got %d", n)
	}
	// A NaN observation would only surface later, rejected by the
	// tail function under its own name; the scan keeps the error with
	// the entry point the caller actually used.
	if err := checkFinite("ChiSquareGoodnessOfFit", "the observed frequencies", observed); err != nil {
		return 0, 0, 0, err
	}
	// A non-finite expectation passes the positivity gate below (+Inf
	// is positive) and turns the statistic into a NaN the tail function
	// would reject under its own, misleading, name.
	if err := checkFinite("ChiSquareGoodnessOfFit", "the expected frequencies", expected); err != nil {
		return 0, 0, 0, err
	}
	chi2 = 0.0
	// Both walks are bounded by the bin count: a rebased view's payload
	// may run past its visible bins, and indexing one past the other's
	// payload would be a panic before it was a wrong statistic.
	obs := rawFloats(observed)
	exp := rawFloats(expected)
	if obs != nil {
		obs = obs[:n]
	}
	if exp != nil {
		exp = exp[:n]
	}
	if obs != nil && exp != nil {
		for i, e := range exp {
			if !(e > 0) {
				return 0, 0, 0, base.Errf("ChiSquareGoodnessOfFit: expected[%d] = %g must be positive", i, e)
			}
			d := obs[i] - e
			chi2 += d * d / e
		}
	} else {
		for i := range n {
			e := expected.FloatAt(i)
			if !(e > 0) {
				return 0, 0, 0, base.Errf("ChiSquareGoodnessOfFit: expected[%d] = %g must be positive", i, e)
			}
			d := observed.FloatAt(i) - e
			chi2 += d * d / e
		}
	}
	df = n - 1
	pValue, err = GammaUpper(float64(df)/2, chi2/2)
	if err != nil {
		return 0, 0, 0, base.Errf("ChiSquareGoodnessOfFit: %w", err)
	}
	return chi2, df, pValue, nil
}

// KolmogorovSmirnovTest compares two samples by the largest vertical
// distance d between their empirical distribution functions. The
// p-value is the asymptotic Kolmogorov distribution evaluated at
// λ = d·√(nm/(n+m)), accurate for samples of a few dozen and upward;
// small samples only get an order-of-magnitude answer. Both samples
// must be finite: a NaN never compares true against the support point
// the merge walk advances on, so it would leave the walk stuck, and an
// ±Inf would distort the distance.
func KolmogorovSmirnovTest(a, b *core.Array) (d, pValue float64, err error) {
	const name = "KolmogorovSmirnovTest"
	if a.Len() == 0 || b.Len() == 0 {
		return 0, 0, base.Errf("%s: both samples must be non-empty", name)
	}
	if a.Dtype() == core.Complex || b.Dtype() == core.Complex {
		return 0, 0, base.Errf("%s: complex samples are not supported", name)
	}
	if err := checkFinite(name, "the first sample", a); err != nil {
		return 0, 0, err
	}
	if err := checkFinite(name, "the second sample", b); err != nil {
		return 0, 0, err
	}
	xs := make([]float64, a.Len())
	if fs := rawFloats(a); fs != nil {
		copy(xs, fs)
	} else {
		for i := range a.Len() {
			xs[i] = a.FloatAt(i)
		}
	}
	ys := make([]float64, b.Len())
	if fs := rawFloats(b); fs != nil {
		copy(ys, fs)
	} else {
		for i := range b.Len() {
			ys[i] = b.FloatAt(i)
		}
	}
	slices.Sort(xs)
	slices.Sort(ys)
	// Walk the merged support, tracking both empirical CDFs.
	i, j := 0, 0
	d = 0.0
	for i < len(xs) || j < len(ys) {
		var x float64
		switch {
		case i < len(xs) && (j >= len(ys) || xs[i] <= ys[j]):
			x = xs[i]
		default:
			x = ys[j]
		}
		for i < len(xs) && xs[i] <= x {
			i++
		}
		for j < len(ys) && ys[j] <= x {
			j++
		}
		if gap := math.Abs(float64(i)/float64(len(xs)) - float64(j)/float64(len(ys))); gap > d {
			d = gap
		}
	}
	en := math.Sqrt(float64(a.Len()) * float64(b.Len()) / float64(a.Len()+b.Len()))
	pValue = kolmogorovTail(d * en)
	return d, pValue, nil
}

// kolmogorovEps is the magnitude at which a term of the Kolmogorov
// series is below the rounding level of the sum, so the alternating
// series can be truncated there: the first omitted term bounds the
// truncation error, and the terms decrease from the first one on.
const kolmogorovEps = 1e-18

// maxKolmogorovTerms caps the series for the smallest lambdas that
// reach it at all.
const maxKolmogorovTerms = 100

// kolmogorovTermCount returns how many terms of the Kolmogorov series
// matter at lambda: the first k whose term 2·exp(−2k²λ²) has fallen
// below the rounding level. The series itself sums exactly that many,
// so the count is the truncation rule in one place instead of a
// condition the loop cannot reach.
func kolmogorovTermCount(lambda float64) int {
	for k := 1; k <= maxKolmogorovTerms; k++ {
		if 2*math.Exp(-2*float64(k)*float64(k)*lambda*lambda) < kolmogorovEps {
			return k
		}
	}
	return maxKolmogorovTerms
}

// kolmogorovTail evaluates the asymptotic Kolmogorov distribution
// Q(λ) = 2·Σ_{k≥1} (−1)^{k−1}·e^{−2k²λ²}, summing the terms that
// matter.
func kolmogorovTail(lambda float64) float64 {
	if lambda < 0.2 {
		return 1
	}
	total := 0.0
	for k := 1; k <= kolmogorovTermCount(lambda); k++ {
		term := 2 * math.Exp(-2*float64(k)*float64(k)*lambda*lambda)
		if k%2 == 0 {
			total -= term
		} else {
			total += term
		}
	}
	return min(1, max(0, total))
}

// BootstrapCI estimates a confidence interval of a statistic by the
// percentile bootstrap: resamples the data with replacement
// resamples times from the seeded generator, evaluates the statistic
// on every resample and returns the alpha/2 and 1−alpha/2 quantiles
// of the bootstrap distribution, alpha = 1 − level. The statistic
// receives the resample as a fresh array and may return its own
// error. The data must be real-valued.
func BootstrapCI(data *core.Array, statistic func(*core.Array) (float64, error), level float64,
	resamples int, seed int64) (lower, upper float64, err error) {
	if data.Len() == 0 {
		return 0, 0, base.Errf("BootstrapCI: the data must not be empty")
	}
	if data.Dtype() == core.Complex {
		return 0, 0, base.Errf("BootstrapCI: complex data are not supported")
	}
	if !(level > 0 && level < 1) {
		return 0, 0, base.Errf("BootstrapCI: level must lie in (0, 1), got %g", level)
	}
	if resamples < 2 {
		return 0, 0, base.Errf("BootstrapCI: resamples must be ≥ 2, got %d", resamples)
	}
	g := core.NewGenerator(seed)
	n := data.Len()
	// The source values are read once and the index payload hoisted:
	// the resample loop then indexes plain slices.
	dataVals := rawFloats(data)
	values := make([]float64, resamples)
	for r := range resamples {
		idx, ierr := core.Ints(g, n, 0, int64(n))
		if ierr != nil {
			return 0, 0, base.Errf("BootstrapCI: %w", ierr)
		}
		idxVals := idx.RawInts()
		// A fresh slice per resample honours the documented contract: a
		// statistic that retains its argument must not observe the next
		// resample's mutation through the alias, and the resample is
		// written into that slice directly rather than into a scratch the
		// caller then copies.
		sample := make([]float64, n)
		if dataVals != nil {
			for i := range n {
				sample[i] = dataVals[idxVals[i]]
			}
		} else {
			for i := range n {
				sample[i] = data.FloatAt(int(idxVals[i]))
			}
		}
		v, serr := statistic(wrapVector(sample))
		if serr != nil {
			return 0, 0, base.Errf("BootstrapCI: %w", serr)
		}
		values[r] = v
	}
	alpha := (1 - level) / 2
	qs, qerr := Quantile(floatsToArray(values, []int{resamples}), []float64{alpha, 1 - alpha})
	if qerr != nil {
		return 0, 0, base.Errf("BootstrapCI: %w", qerr)
	}
	return qs.FloatAt(0), qs.FloatAt(1), nil
}

// wrapVector views a float64 slice as a rank-1 Array without copying.
func wrapVector(v []float64) *core.Array {
	a, _ := core.FloatsFromArray(v, len(v))
	return a
}