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

package signal

import (
	"math"
	"testing"

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

// arSeries generates a zero-mean AR(p) series: x_t = Σ φ_j·x_{t−j} +
// e_t with standard normal innovations, the first burn values of the
// burn-in discarded.
func arSeries(g *core.Generator, phi []float64, n, burn int) []float64 {
	p := len(phi)
	x := make([]float64, n+burn)
	for i := p; i < len(x); i++ {
		v := g.NormalUnit()
		for j := range p {
			v += phi[j] * x[i-1-j]
		}
		x[i] = v
	}
	return x[burn:]
}

// arma11Series generates a zero-mean ARMA(1,1) series: x_t =
// φ·x_{t−1} + e_t + θ·e_{t−1}, the burn-in discarded.
func arma11Series(g *core.Generator, phi, theta float64, n, burn int) []float64 {
	x := make([]float64, n)
	var xPrev, ePrev float64
	for i := -burn; i < n; i++ {
		e := g.NormalUnit()
		v := phi*xPrev + e + theta*ePrev
		xPrev, ePrev = v, e
		if i >= 0 {
			x[i] = v
		}
	}
	return x
}

// TestEstimateAR2YuleWalker recovers a known AR(2) from a long
// generated series, and then demands the recovered coefficients solve
// the Yule-Walker equations on the empirical autocovariance the fit
// itself read: the Durbin-Levinson solve is exact, so the residuals
// sit at rounding level. The partial autocorrelation, produced by the
// same recursion answering a different question, must agree: lag one
// carries φ1, lag two the signature φ2/(1 − φ1²)... for an AR(2) the
// second partial is φ2 and the rest are noise.
func TestEstimateAR2YuleWalker(t *testing.T) {
	const (
		phi1 = 0.8
		phi2 = -0.4
		n    = 60000
	)
	g := core.NewGenerator(7)
	x := arSeries(g, []float64{phi1, phi2}, n, 500)
	res, err := EstimateAR(mustFloats(t, x), 2)
	if err != nil {
		t.Fatalf("EstimateAR: %v", err)
	}
	if math.Abs(res.AR[0]-phi1) > 0.02 || math.Abs(res.AR[1]-phi2) > 0.02 {
		t.Fatalf("AR(2) coefficients (%.4f, %.4f), want (%.2f, %.2f)", res.AR[0], res.AR[1], phi1, phi2)
	}
	if math.Abs(res.InnovationVariance-1) > 0.1 {
		t.Fatalf("innovation variance %.4f, want 1", res.InnovationVariance)
	}
	if math.Abs(res.Mean) > 0.05 {
		t.Fatalf("mean %.4f, want 0", res.Mean)
	}
	// The Yule-Walker equations on the empirical autocovariance.
	r, _, _, err := armaAutocovariance("test", mustFloats(t, x), 2)
	if err != nil {
		t.Fatalf("armaAutocovariance: %v", err)
	}
	for k := range 2 {
		got := res.AR[k]*r[0] + res.AR[1-k]*r[1] - r[k+1]
		if math.Abs(got) > 1e-8*r[0] {
			t.Fatalf("Yule-Walker residual at lag %d = %g, want rounding only", k+1, got)
		}
	}
	// The PACF, the same recursion read for its reflection
	// coefficients, must agree with the fit: lag two's partial is the
	// order-two coefficient itself, lag one's is the lag-one
	// autocorrelation φ1/(1 − φ2), both exact relations of the same
	// autocovariances the Yule-Walker solve consumed.
	pacf, err := PartialAutocorrelate(mustFloats(t, x), 3)
	if err != nil {
		t.Fatalf("PartialAutocorrelate: %v", err)
	}
	if math.Abs(pacf.FloatAt(1)-res.AR[1]) > 1e-9 {
		t.Fatalf("PACF lag 2 = %.10f against the fitted φ2 %.10f", pacf.FloatAt(1), res.AR[1])
	}
	rho1 := res.AR[0] / (1 - res.AR[1])
	if math.Abs(pacf.FloatAt(0)-rho1) > 1e-9 {
		t.Fatalf("PACF lag 1 = %.10f against φ1/(1 − φ2) = %.10f", pacf.FloatAt(0), rho1)
	}
	if math.Abs(pacf.FloatAt(2)) > 0.02 {
		t.Fatalf("PACF lag 3 = %.4f, an AR(2) cuts off after lag 2", pacf.FloatAt(2))
	}
	// The criteria carry the concentrated likelihood of the whole
	// series with one parameter per coefficient plus the variance.
	wantAIC := float64(n)*(math.Log(2*math.Pi)+math.Log(res.InnovationVariance)+1) + 2*3
	if math.Abs(res.AIC-wantAIC) > 1e-6 {
		t.Fatalf("AIC %.6f, want %.6f", res.AIC, wantAIC)
	}
	wantBIC := float64(n)*(math.Log(2*math.Pi)+math.Log(res.InnovationVariance)+1) + math.Log(float64(n))*3
	if math.Abs(res.BIC-wantBIC) > 1e-6 {
		t.Fatalf("BIC %.6f, want %.6f", res.BIC, wantBIC)
	}
}

// TestEstimateARErrors pins the input gates and the stationary-data
// requirement.
func TestEstimateARErrors(t *testing.T) {
	x, _ := core.FromFloats([]float64{1, 2, 3, 4, 5, 6}, 6)
	if _, err := EstimateAR(x, 0); err == nil {
		t.Error("order zero accepted")
	}
	if _, err := EstimateAR(x, 5); err == nil {
		t.Error("order at the sample count accepted")
	}
	rank2, _ := core.FromFloats([]float64{1, 2, 3, 4}, 2, 2)
	if _, err := EstimateAR(rank2, 1); err == nil {
		t.Error("rank-2 series accepted")
	}
	bad, _ := core.FromFloats([]float64{1, math.NaN(), 3, 4}, 4)
	if _, err := EstimateAR(bad, 1); err == nil {
		t.Error("non-finite series accepted")
	}
	constant, _ := core.FromFloats([]float64{2, 2, 2, 2, 2, 2}, 6)
	if _, err := EstimateAR(constant, 1); err == nil {
		t.Error("constant series accepted")
	}
}

// TestEstimateARMA11HannanRissanen recovers a known ARMA(1,1) from a
// long generated series. The honest caveat, documented with the
// function: Hannan-Rissen conditions on proxy residuals, so the
// estimates carry more sampling noise than a maximum-likelihood fit
// would and the tolerances here are an order looser than the AR(2)
// ones above, the MA side most of all.
func TestEstimateARMA11HannanRissanen(t *testing.T) {
	const (
		phi   = 0.6
		theta = 0.4
		n     = 80000
	)
	g := core.NewGenerator(11)
	x := arma11Series(g, phi, theta, n, 500)
	res, err := EstimateARMA(mustFloats(t, x), 1, 1, ARMAOptions{})
	if err != nil {
		t.Fatalf("EstimateARMA: %v", err)
	}
	if res.P != 1 || res.Q != 1 {
		t.Fatalf("orders (%d, %d), want (1, 1)", res.P, res.Q)
	}
	if math.Abs(res.AR[0]-phi) > 0.05 {
		t.Fatalf("AR coefficient %.4f, want %.2f within 0.05", res.AR[0], phi)
	}
	if math.Abs(res.MA[0]-theta) > 0.05 {
		t.Fatalf("MA coefficient %.4f, want %.2f within 0.05", res.MA[0], theta)
	}
	if math.Abs(res.InnovationVariance-1) > 0.15 {
		t.Fatalf("innovation variance %.4f, want 1 within 0.15", res.InnovationVariance)
	}
	if math.Abs(res.Mean) > 0.05 {
		t.Fatalf("mean %.4f, want 0", res.Mean)
	}
	// The criteria carry the concentrated whole-series likelihood
	// here too, not one windowed by the high-order AR proxy: a
	// regression to the windowed count shifts AIC by dozens of points
	// and breaks cross-model comparability, so the closed form pins
	// the value on the ARMA path exactly as it does on the AR one.
	wantLL := -float64(n) / 2 * (math.Log(2*math.Pi) + math.Log(res.InnovationVariance) + 1)
	if math.Abs(res.LogLikelihood-wantLL) > 1e-6 {
		t.Fatalf("log likelihood %.6f, want %.6f", res.LogLikelihood, wantLL)
	}
	wantAIC := -2*res.LogLikelihood + 2*3
	if math.Abs(res.AIC-wantAIC) > 1e-6 {
		t.Fatalf("AIC %.6f, want %.6f", res.AIC, wantAIC)
	}
}

// TestEstimateARMAErrors pins the order and data gates, including the
// routing of a pure AR model to the Yule-Walker fit.
func TestEstimateARMAErrors(t *testing.T) {
	x, _ := core.FromFloats([]float64{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, 10)
	if _, err := EstimateARMA(x, 0, 0, ARMAOptions{}); err == nil {
		t.Error("orders (0, 0) accepted")
	}
	if _, err := EstimateARMA(x, -1, 1, ARMAOptions{}); err == nil {
		t.Error("negative order accepted")
	}
	if _, err := EstimateARMA(x, 1, 1, ARMAOptions{HighAROrder: 1}); err == nil {
		t.Error("proxy order below p+q+1 accepted")
	}
	if _, err := EstimateARMA(x, 1, 1, ARMAOptions{}); err == nil {
		t.Error("series too short for the regression window accepted")
	}
	// A pure AR request is routed to the Yule-Walker fit: orders below
	// its own gate are refused there.
	if _, err := EstimateARMA(x, 0, 1, ARMAOptions{}); err == nil {
		t.Error("degenerate ARMA with p = 0 accepted on a short series")
	}
	constant, _ := core.FromFloats([]float64{2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2}, 20)
	if _, err := EstimateARMA(constant, 1, 1, ARMAOptions{}); err == nil {
		t.Error("constant series accepted")
	}
}

// TestSelectARMAInformationCriteria generates an AR(2) and demands the
// grid search pick exactly (2, 0) by both criteria, then an ARMA(1,1)
// and demands both criteria pick exactly (1, 1). Both pins are
// deterministic on the seeded generator: the streams are bit-stable,
// so the choice is a regression pin, not a restatement of the
// criteria's large-sample guarantees. The seeds are chosen for wide
// decision margins (the closest competitor sits 1.999 AIC points
// behind on the AR(2) grid, 1.991 AIC and 11.3 BIC points on the
// ARMA(1,1) grid), because the AIC's documented willingness to
// overfit by one order makes exact selection a coin weighted by the
// sample: with margin 2 the decision cannot move on arithmetic-level
// perturbations.
func TestSelectARMAInformationCriteria(t *testing.T) {
	g := core.NewGenerator(29)
	x := arSeries(g, []float64{0.8, -0.4}, 60000, 500)
	for _, criterion := range []string{"aic", "bic"} {
		res, err := SelectARMA(mustFloats(t, x), 3, 1, ARMAOptions{Criterion: criterion})
		if err != nil {
			t.Fatalf("SelectARMA(%s): %v", criterion, err)
		}
		if res.P != 2 || res.Q != 0 {
			t.Fatalf("%s picked ARMA(%d, %d), want ARMA(2, 0)", criterion, res.P, res.Q)
		}
	}
	// Both criteria of the winning model are reported.
	res, err := SelectARMA(mustFloats(t, x), 3, 1, ARMAOptions{Criterion: "bic"})
	if err != nil {
		t.Fatalf("SelectARMA: %v", err)
	}
	if !(res.BIC > res.AIC) {
		t.Fatalf("BIC %.4f does not exceed AIC %.4f at n = 60000", res.BIC, res.AIC)
	}
	// The ARMA(1,1) series on a wider grid: both criteria pin (1, 1).
	g2 := core.NewGenerator(4)
	y := arma11Series(g2, 0.6, 0.4, 80000, 500)
	for _, criterion := range []string{"aic", "bic"} {
		best, err := SelectARMA(mustFloats(t, y), 2, 2, ARMAOptions{Criterion: criterion})
		if err != nil {
			t.Fatalf("SelectARMA(%s): %v", criterion, err)
		}
		if best.P != 1 || best.Q != 1 {
			t.Fatalf("%s picked ARMA(%d, %d), want ARMA(1, 1)", criterion, best.P, best.Q)
		}
	}
}

// TestSelectARMAErrors pins the grid and criterion gates, including
// the reporting when no cell can be fitted.
func TestSelectARMAErrors(t *testing.T) {
	x, _ := core.FromFloats([]float64{1, 2, 3, 4, 5, 6}, 6)
	if _, err := SelectARMA(x, 0, 0, ARMAOptions{}); err == nil {
		t.Error("an empty grid accepted")
	}
	if _, err := SelectARMA(x, -1, 2, ARMAOptions{}); err == nil {
		t.Error("a negative grid bound accepted")
	}
	if _, err := SelectARMA(x, 1, 1, ARMAOptions{Criterion: "sic"}); err == nil {
		t.Error("an unknown criterion accepted")
	}
	constant, _ := core.FromFloats([]float64{2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2}, 18)
	if _, err := SelectARMA(constant, 2, 2, ARMAOptions{}); err == nil {
		t.Error("a grid where nothing can be fitted reported no error")
	}
}

// TestARMASpectrumMatchesPeriodogram evaluates the fitted AR(2) model's
// theoretical spectrum against Welch's estimate of a long series
// generated from that fitted model: the frequency grids line up bin
// for bin, and on the broad band between 0.03 and 0.47
// cycles-per-sample the two log-spectra must track each other well
// inside the averaging fluctuation of the estimate.
func TestARMASpectrumMatchesPeriodogram(t *testing.T) {
	const n = 40000
	g := core.NewGenerator(31)
	x := arSeries(g, []float64{0.8, -0.4}, n, 500)
	fitted, err := EstimateAR(mustFloats(t, x), 2)
	if err != nil {
		t.Fatalf("EstimateAR: %v", err)
	}
	// A long series from the fitted coefficients: the fitted model's
	// own output, whose periodogram the theory must describe.
	const n2 = 1 << 18
	sigma := math.Sqrt(fitted.InnovationVariance)
	y := make([]float64, n2+500)
	for i := 2; i < len(y); i++ {
		y[i] = fitted.AR[0]*y[i-1] + fitted.AR[1]*y[i-2] + sigma*g.NormalUnit()
	}
	const segment = 1024
	freqs, psd, err := WelchPSD(mustFloats(t, y[500:]), 1, segment, segment/2, "hann")
	if err != nil {
		t.Fatalf("WelchPSD: %v", err)
	}
	bins := segment/2 + 1
	tFreqs, tPsd, err := ARMASpectrum(fitted, bins)
	if err != nil {
		t.Fatalf("ARMASpectrum: %v", err)
	}
	for k := range bins {
		if math.Abs(freqs.FloatAt(k)-tFreqs.FloatAt(k)) > 1e-12 {
			t.Fatalf("frequency grids disagree at bin %d", k)
		}
	}
	worst := 0.0
	for k := range bins {
		f := freqs.FloatAt(k)
		if f < 0.03 || f > 0.47 {
			continue
		}
		d := math.Abs(math.Log(psd.FloatAt(k)) - math.Log(tPsd.FloatAt(k)))
		if d > worst {
			worst = d
		}
	}
	if worst > 0.2 {
		t.Fatalf("log-spectral gap %.4f on the broad band, want under 0.2", worst)
	}
	// The resonance: cos ω₀ = φ1(φ2 − 1)/(4φ2) locates the AR(2)
	// peak; both the Welch estimate and the theoretical spectrum must
	// peak inside ±0.03 cycles-per-sample of it, and the theoretical
	// argmax must sit within two bins of the formula's own optimum
	// (the peak is broad, so bins are the honest resolution).
	phi1, phi2 := fitted.AR[0], fitted.AR[1]
	cosw0 := phi1 * (phi2 - 1) / (4 * phi2)
	f0 := math.Acos(cosw0) / (2 * math.Pi)
	if f0 < 0.02 || f0 > 0.48 {
		t.Fatalf("the resonance formula left the band: %g", f0)
	}
	binOf := func(a *core.Array) int {
		best, at := math.Inf(-1), -1
		for k := range bins {
			f := freqs.FloatAt(k)
			if f < 0.02 || f > 0.48 {
				continue
			}
			if a.FloatAt(k) > best {
				best, at = a.FloatAt(k), k
			}
		}
		return at
	}
	if d := math.Abs(float64(binOf(psd))/float64(segment) - f0); d > 0.03 {
		t.Fatalf("the empirical peak sits %.4f from the resonance %g", d, f0)
	}
	if d := math.Abs(float64(binOf(tPsd))/float64(segment) - f0); d > 2.0/float64(segment) {
		t.Fatalf("the theoretical peak sits %.4f from the resonance %g", d, f0)
	}
	// The theoretical values, checked against the formula evaluated
	// independently on a fine grid: the ARMA(1,0) polynomial arithmetic
	// is the whole function, so this pins it to rounding.
	const fine = 4097
	fFreqs, fPsd, err := ARMASpectrum(fitted, fine)
	if err != nil {
		t.Fatalf("ARMASpectrum: %v", err)
	}
	for k := range fine {
		f := 0.5 * float64(k) / float64(fine-1)
		w := 2 * math.Pi * f
		re := 1 - phi1*math.Cos(w) - phi2*math.Cos(2*w)
		im := phi1*math.Sin(w) + phi2*math.Sin(2*w)
		want := fitted.InnovationVariance / (re*re + im*im)
		if k > 0 && k < fine-1 {
			want *= 2
		}
		if math.Abs(fPsd.FloatAt(k)-want) > 1e-10*want {
			t.Fatalf("spectrum %.12g at bin %d, formula %.12g", fPsd.FloatAt(k), k, want)
		}
		if math.Abs(fFreqs.FloatAt(k)-f) > 1e-12 {
			t.Fatalf("frequency %.12g at bin %d, want %.12g", fFreqs.FloatAt(k), k, f)
		}
	}
	// The flat model answers the innovation variance at the edges and
	// its double between them: the one-sided fold the Welch estimate
	// of the same white noise prints.
	flat := &ARMAResult{InnovationVariance: 2.5}
	_, flatPsd, err := ARMASpectrum(flat, 16)
	if err != nil {
		t.Fatalf("ARMASpectrum: %v", err)
	}
	for k := range 16 {
		want := 5.0
		if k == 0 || k == 15 {
			want = 2.5
		}
		if flatPsd.FloatAt(k) != want {
			t.Fatalf("the flat model's spectrum is %g at bin %d, want %g", flatPsd.FloatAt(k), k, want)
		}
	}
}

// TestARMASpectrumErrors pins the gates, including the pole on the
// unit circle being named instead of answered with infinities.
func TestARMASpectrumErrors(t *testing.T) {
	if _, _, err := ARMASpectrum(nil, 8); err == nil {
		t.Error("a nil model accepted")
	}
	res := &ARMAResult{AR: []float64{0.5}, InnovationVariance: 1}
	if _, _, err := ARMASpectrum(res, 1); err == nil {
		t.Error("a single frequency accepted")
	}
	res.InnovationVariance = -1
	if _, _, err := ARMASpectrum(res, 8); err == nil {
		t.Error("a negative innovation variance accepted")
	}
	res.InnovationVariance = 1
	res.AR = []float64{math.NaN()}
	if _, _, err := ARMASpectrum(res, 8); err == nil {
		t.Error("a non-finite coefficient accepted")
	}
	res.AR = []float64{1}
	if _, _, err := ARMASpectrum(res, 8); err == nil {
		t.Error("a pole on the unit circle accepted")
	}
}