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

package signal

import (
	"math"

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

// Time-series model estimation: the autoregressive and
// ARMA staples. The pure AR fit solves the Yule-Walker equations by
// the Durbin-Levinson recursion over the biased autocovariance, the
// small solver this package's own PartialAutocorrelate runs in
// miniature (the recursion lives locally here; the domain boundary
// allows no linalg import, and the earlier recursion answers a
// different question). The mixed model is fitted by the
// Hannan-Rissanen innovations method: a high-order autoregression
// first, whose residuals stand in for the unobserved innovations,
// then one least-squares regression of the series on its own lags and
// the lagged residuals. Both report the concentrated Gaussian
// likelihood of their innovations, and with it the AIC and BIC the
// grid search in SelectARMA minimises.

// ARMAResult holds one fitted time-series model. The coefficients
// follow the convention x_t − μ = Σ_j φ_j·(x_{t−j} − μ) + e_t +
// Σ_k θ_k·e_{t−k}, innovations white with variance
// InnovationVariance.
type ARMAResult struct {
	// AR holds φ_1..φ_P, empty only for a pure MA model.
	AR []float64
	// MA holds θ_1..θ_Q, empty only for a pure AR model.
	MA []float64
	// InnovationVariance is the estimated variance σ² of e_t.
	InnovationVariance float64
	// Mean is the sample mean μ the fit removed and reports back.
	Mean float64
	// LogLikelihood is the concentrated Gaussian likelihood of the
	// innovations over the whole series: −n/2·(log 2πσ² + 1) at the
	// fitted innovation variance, so the criteria of different models
	// on one series sit on a common footing.
	LogLikelihood float64
	// AIC is −2·LogLikelihood + 2k with k = P + Q + 1, the variance
	// counted as a parameter.
	AIC float64
	// BIC is −2·LogLikelihood + k·log n, the same k against the
	// innovation count's logarithm.
	BIC float64
	// P and Q record the orders the result carries.
	P int
	Q int
}

// ARMAOptions tunes the Hannan-Rissanen estimation and the grid
// search. HighAROrder is the order of the proxy autoregression whose
// residuals supply the MA stage; the zero value implies the default
// max(16, p+q+8), which the data length bounds. Criterion names the
// information criterion SelectARMA minimises: "aic" or "bic", with
// "aic" implied by the zero value.
type ARMAOptions struct {
	HighAROrder int
	Criterion   string
}

// armaLikelihood fills the concentrated likelihood and the criteria
// from the innovation variance, evaluated over the whole series: nEff
// is the same sample count for every model fitted to one series,
// which is what makes the criteria comparable across a grid, and k
// counts the coefficients plus the variance. The values a fit
// conditions on (an AR's first p samples, the Hannan-Rissanen
// window's proxy residuals) are estimation detail, not a data
// difference.
func armaLikelihood(res *ARMAResult, nEff int) {
	k := res.P + res.Q + 1
	logL := -0.5 * float64(nEff) * (math.Log(2*math.Pi) + math.Log(res.InnovationVariance) + 1)
	res.LogLikelihood = logL
	res.AIC = 2*float64(k) - 2*logL
	res.BIC = math.Log(float64(nEff))*float64(k) - 2*logL
}

// armaAutocovariance returns the biased autocovariances r_0..r_maxLag
// of a real rank-1 series, the estimator the Yule-Walker theory is
// written for: the mean removed, every lag divided by the sample
// count. The series' values are read through the house autocorrelation
// and rescaled by the lag-zero power. The mean comes back too, because
// every model is fitted to the demeaned series and reported with it.
func armaAutocovariance(name string, x *core.Array, maxLag int) (r []float64, n int, mean float64, err error) {
	if x.NDim() != 1 {
		return nil, 0, 0, base.Errf("%s: needs a rank-1 series, got shape %s", name, base.ShapeText(x.Shape()))
	}
	if x.Dtype() == core.Complex {
		return nil, 0, 0, base.Errf("%s: complex series are not supported", name)
	}
	n = x.Len()
	if n < maxLag+2 {
		return nil, 0, 0, base.Errf("%s: at least %d samples are needed for lag %d, got %d",
			name, maxLag+2, maxLag, n)
	}
	vals := widenFloats(x)
	if err := kfFinite(name, "the series", vals); err != nil {
		return nil, 0, 0, err
	}
	for _, v := range vals {
		mean += v
	}
	mean /= float64(n)
	power := 0.0
	for _, v := range vals {
		d := v - mean
		power += d * d
	}
	variance := power / float64(n)
	acf, err := Autocorrelate(x, maxLag)
	if err != nil {
		return nil, 0, 0, err
	}
	r = make([]float64, maxLag+1)
	for k := range r {
		r[k] = acf.FloatAt(k) * variance
	}
	return r, n, mean, nil
}

// yuleWalker solves the Yule-Walker equations R·φ = r for the
// order-p autoregression's coefficients, R the Toeplitz matrix of the
// autocovariances r[0..p], through the Durbin-Levinson recursion: p²
// work, no matrix factored. It returns φ_1..φ_p and the innovation
// variance r₀·(1 − Σφ_j·ρ_j). A non-positive recursion denominator
// means the autocovariances do not belong to a stationary process;
// the same recursion, run for its reflection coefficients, is what
// PartialAutocorrelate reports.
func yuleWalker(name string, r []float64, p int) (phi []float64, sigma2 float64, err error) {
	phi = make([]float64, p+1) // 1-based: phi[k] holds φ_k
	prev := make([]float64, p+1)
	for k := 1; k <= p; k++ {
		num, den := r[k], r[0]
		for j := 1; j < k; j++ {
			num -= prev[j] * r[k-j]
			den -= prev[j] * r[j]
		}
		if !(den > 0) || math.IsNaN(den) {
			return nil, 0, base.Errf("%s: the Durbin-Levinson recursion broke down at order %d: the autocovariances do not describe a stationary process", name, k)
		}
		phi[k] = num / den
		for j := 1; j < k; j++ {
			phi[j] = prev[j] - phi[k]*prev[k-j]
		}
		copy(prev, phi)
	}
	out := make([]float64, p)
	copy(out, phi[1:])
	sigma2 = r[0]
	for j := 1; j <= p; j++ {
		sigma2 -= phi[j] * r[j]
	}
	if !(sigma2 > 0) || math.IsNaN(sigma2) {
		return nil, 0, base.Errf("%s: the innovation variance came out %g: the autocovariances do not describe a stationary process", name, sigma2)
	}
	return out, sigma2, nil
}

// armaLstsq solves the small least-squares problem min ‖A·c − b‖²
// through the normal equations and Gaussian elimination with partial
// pivoting. The width stays at a model order's size, where the
// squandered conditioning costs nothing that matters; a pivot too
// small against the normal matrix's largest entry reports a
// degenerate regression rather than returning coefficients for it.
func armaLstsq(name string, design [][]float64, target []float64) ([]float64, error) {
	width := len(design[0])
	g := make([]float64, width*width)
	rhs := make([]float64, width)
	for i, row := range design {
		for a := range width {
			rhs[a] += row[a] * target[i]
			for b := range width {
				g[a*width+b] += row[a] * row[b]
			}
		}
	}
	big := 0.0
	for _, v := range g {
		big = max(big, math.Abs(v))
	}
	for a := range width {
		piv := a
		for c := a + 1; c < width; c++ {
			if math.Abs(g[c*width+a]) > math.Abs(g[piv*width+a]) {
				piv = c
			}
		}
		if !(math.Abs(g[piv*width+a]) > 1e-12*big) {
			return nil, base.Errf("%s: the regression design is degenerate: its normal equations have no pivot at column %d", name, a+1)
		}
		if piv != a {
			for b := range width {
				g[a*width+b], g[piv*width+b] = g[piv*width+b], g[a*width+b]
			}
			rhs[a], rhs[piv] = rhs[piv], rhs[a]
		}
		for c := a + 1; c < width; c++ {
			f := g[c*width+a] / g[a*width+a]
			for b := a; b < width; b++ {
				g[c*width+b] -= f * g[a*width+b]
			}
			rhs[c] -= f * rhs[a]
		}
	}
	coef := make([]float64, width)
	for a := width - 1; a >= 0; a-- {
		total := rhs[a]
		for b := a + 1; b < width; b++ {
			total -= g[a*width+b] * coef[b]
		}
		coef[a] = total / g[a*width+a]
	}
	return coef, nil
}

// EstimateAR fits the order-p autoregression x_t − μ = Σ_j
// φ_j·(x_{t−j} − μ) + e_t to the real rank-1 series x: the
// Yule-Walker equations over the biased autocovariance, solved by the
// Durbin-Levinson recursion. The result carries the coefficients, the
// innovation variance, the concentrated Gaussian likelihood of the
// whole series, and the AIC and BIC that compare models fitted to
// one series. An order below 1, a series too short for the lag, and
// non-finite or complex data are errors; autocovariances that do not
// describe a stationary process break the recursion and are reported.
func EstimateAR(x *core.Array, order int) (*ARMAResult, error) {
	const name = "EstimateAR"
	if order < 1 {
		return nil, base.Errf("%s: the order must be at least 1, got %d", name, order)
	}
	r, n, mean, err := armaAutocovariance(name, x, order)
	if err != nil {
		return nil, err
	}
	phi, sigma2, err := yuleWalker(name, r, order)
	if err != nil {
		return nil, err
	}
	res := &ARMAResult{AR: phi, InnovationVariance: sigma2, Mean: mean, P: order}
	armaLikelihood(res, n)
	return res, nil
}

// EstimateARMA fits the ARMA(p, q) model x_t − μ = Σ_j φ_j·(x_{t−j} −
// μ) + e_t + Σ_k θ_k·e_{t−k} by the Hannan-Rissanen innovations
// method: an order-m autoregression first (m from ARMAOptions'
// HighAROrder, defaulting to max(16, p+q+8)), whose residuals stand
// in for the unobserved innovations, then one least-squares
// regression of the demeaned series on its own lags and the lagged
// residuals. The result carries the coefficients in the documented
// convention, the regression residuals' variance as the innovation
// variance and the likelihood and criteria computed from it over the
// whole series, so the criteria line up with the pure AR fit's and
// with the other cells of a SelectARMA grid.
//
// The honest caveat: the method conditions on proxy residuals, so its
// estimates carry more sampling noise than a maximum-likelihood fit
// would, the MA side most of all; tolerances set against it should be
// correspondingly looser. A pure AR model (q = 0) is routed to the
// Yule-Walker fit, which solves that problem exactly; p + q must be
// at least 1, the proxy order at least p+q+1, and the series long
// enough to leave a regression window worth fitting.
func EstimateARMA(x *core.Array, p, q int, opts ARMAOptions) (*ARMAResult, error) {
	const name = "EstimateARMA"
	if p < 0 || q < 0 {
		return nil, base.Errf("%s: the orders must not be negative, got %d and %d", name, p, q)
	}
	if p+q < 1 {
		return nil, base.Errf("%s: at least one coefficient is needed, got orders %d and %d", name, p, q)
	}
	if q == 0 {
		return EstimateAR(x, p)
	}
	m := opts.HighAROrder
	if m == 0 {
		m = max(16, p+q+8)
	}
	if m < p+q+1 {
		return nil, base.Errf("%s: the proxy AR order must be at least p+q+1 = %d, got %d", name, p+q+1, m)
	}
	r, n, mean, err := armaAutocovariance(name, x, m)
	if err != nil {
		return nil, err
	}
	rows := n - m - q
	if rows < p+q+8 {
		return nil, base.Errf("%s: %d samples leave only %d regression rows for %d coefficients; shorten the proxy or lengthen the series",
			name, n, rows, p+q)
	}
	phiAR, _, err := yuleWalker(name, r[:m+1], m)
	if err != nil {
		return nil, err
	}
	demeaned := make([]float64, n)
	for i := range n {
		demeaned[i] = x.FloatAt(i) - mean
	}
	// The proxy AR's residuals, the innovations' stand-ins: e_t for
	// t ≥ m.
	resid := make([]float64, n-m)
	for t := m; t < n; t++ {
		v := demeaned[t]
		for j := 1; j <= m; j++ {
			v -= phiAR[j-1] * demeaned[t-j]
		}
		resid[t-m] = v
	}
	// The regression: x_t − μ on x_{t−1..t−p} − μ and e_{t−1..t−q},
	// over the window t = m+q..n−1 the proxies cover.
	T := n - m - q
	design := make([][]float64, T)
	target := make([]float64, T)
	for i := range T {
		t := m + q + i
		row := make([]float64, 0, p+q)
		for j := 1; j <= p; j++ {
			row = append(row, demeaned[t-j])
		}
		for k := 1; k <= q; k++ {
			row = append(row, resid[t-k-m])
		}
		design[i] = row
		target[i] = demeaned[t]
	}
	coef, err := armaLstsq(name, design, target)
	if err != nil {
		return nil, err
	}
	rss := 0.0
	for i := range T {
		fitted := 0.0
		for a, v := range design[i] {
			fitted += v * coef[a]
		}
		d := target[i] - fitted
		rss += d * d
	}
	sigma2 := rss / float64(T)
	if !(sigma2 > 0) || math.IsNaN(sigma2) {
		return nil, base.Errf("%s: the innovation variance came out %g: the fit carries no information", name, sigma2)
	}
	res := &ARMAResult{AR: coef[:p], MA: coef[p:], InnovationVariance: sigma2, Mean: mean, P: p, Q: q}
	armaLikelihood(res, n)
	return res, nil
}

// SelectARMA searches the order grid 0..maxAR × 0..maxMA for the ARMA
// model the chosen information criterion ranks best: every cell is
// fitted by EstimateARMA, the criterion (ARMAOptions' Criterion,
// "aic" by default) is computed from the innovation variance, and the
// smallest score wins. The (0, 0) cell carries no dynamics and is
// skipped; a cell whose fit fails is skipped too, and a grid where
// nothing can be fitted reports the last failure. The result carries
// both criteria of the winning model, so the runner-up is one more
// sweep away.
//
// The two criteria answer different questions: the AIC's linear
// parameter penalty makes it efficient but willing to overfit by an
// order with probability that does not vanish with the sample count,
// while the BIC's logarithmic penalty makes it consistent, choosing
// the true orders with probability tending to one. Where they
// disagree on long series, the BIC is the one to trust.
func SelectARMA(x *core.Array, maxAR, maxMA int, opts ARMAOptions) (*ARMAResult, error) {
	const name = "SelectARMA"
	if maxAR < 0 || maxMA < 0 {
		return nil, base.Errf("%s: the grid bounds must not be negative, got %d and %d", name, maxAR, maxMA)
	}
	if maxAR+maxMA < 1 {
		return nil, base.Errf("%s: the grid must hold at least one candidate, got %d×%d", name, maxAR+1, maxMA+1)
	}
	criterion := opts.Criterion
	if criterion == "" {
		criterion = "aic"
	}
	if criterion != "aic" && criterion != "bic" {
		return nil, base.Errf("%s: the criterion must be %q or %q, got %q", name, "aic", "bic", opts.Criterion)
	}
	var (
		best      *ARMAResult
		bestScore float64
		lastErr   error
	)
	for p := range maxAR + 1 {
		for q := range maxMA + 1 {
			if p == 0 && q == 0 {
				continue
			}
			res, err := EstimateARMA(x, p, q, opts)
			if err != nil {
				lastErr = err
				continue
			}
			score := res.AIC
			if criterion == "bic" {
				score = res.BIC
			}
			if best == nil || score < bestScore {
				best, bestScore = res, score
			}
		}
	}
	if best == nil {
		return nil, base.Errf("%s: no model on the grid could be fitted: %w", name, lastErr)
	}
	return best, nil
}

// ARMASpectrum evaluates the fitted model's theoretical power
// spectrum at the nFreq frequencies evenly spaced from 0 to the
// Nyquist frequency 0.5 inclusive, on the one-sided convention the
// house periodogram prints: the bins strictly between DC and Nyquist
// carry the folded double weight, DC and Nyquist the single one. The
// value is σ²·|Θ(e^{−2πif})|² / |Φ(e^{−2πif})|² with that weighting,
// so model and data line up bin for bin against WelchPSD at fs = 1:
// a white-noise sequence of variance σ² estimates σ² at the edges and
// 2σ² between them, and the fitted model's spectrum describes the
// data's own periodogram without a scale fudge. Fewer than two
// frequencies, a nil model, and a non-positive innovation variance or
// non-finite coefficient are errors; a pole of Φ on the unit circle
// names its frequency instead of answering infinities.
func ARMASpectrum(res *ARMAResult, nFreq int) (freqs, psd *core.Array, err error) {
	const name = "ARMASpectrum"
	if res == nil {
		return nil, nil, base.Errf("%s: the model is nil", name)
	}
	if nFreq < 2 {
		return nil, nil, base.Errf("%s: at least two frequencies are needed, got %d", name, nFreq)
	}
	if !(res.InnovationVariance > 0) || math.IsNaN(res.InnovationVariance) || math.IsInf(res.InnovationVariance, 0) {
		return nil, nil, base.Errf("%s: the innovation variance must be positive and finite, got %g",
			name, res.InnovationVariance)
	}
	vals := append(append([]float64{}, res.AR...), res.MA...)
	if err := kfFinite(name, "the coefficients", vals); err != nil {
		return nil, nil, err
	}
	freqs = core.New(core.Float, nFreq)
	psd = core.New(core.Float, nFreq)
	fRow := freqs.RawFloats()
	pRow := psd.RawFloats()
	for k := range nFreq {
		f := 0.5 * float64(k) / float64(nFreq-1)
		omega := 2 * math.Pi * f
		// |Φ|² and |Θ|² by real arithmetic: the phase parts multiply
		// out to squares of sine sums. Both loops read a sine and a
		// cosine of the same argument; math.Sincos returns exactly the
		// pair math.Sin and math.Cos produce (verified bit-for-bit), so
		// the sums are unchanged while the trig work halves.
		rePhi, imPhi := 1.0, 0.0
		for j, phi := range res.AR {
			s, c := math.Sincos(float64(j+1) * omega)
			rePhi -= phi * c
			imPhi += phi * s
		}
		reTheta, imTheta := 1.0, 0.0
		for k2, theta := range res.MA {
			s, c := math.Sincos(float64(k2+1) * omega)
			reTheta += theta * c
			imTheta -= theta * s
		}
		den := rePhi*rePhi + imPhi*imPhi
		if den == 0 {
			return nil, nil, base.Errf("%s: the model has a pole on the unit circle at frequency %g", name, f)
		}
		fRow[k] = f
		pRow[k] = res.InnovationVariance * (reTheta*reTheta + imTheta*imTheta) / den
		if k > 0 && k < nFreq-1 {
			pRow[k] *= 2 // the one-sided fold between DC and Nyquist
		}
	}
	return freqs, psd, nil
}