// 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"
)

// Correlation functions: the time-series staples on top of
// the FFT. The autocorrelation and cross-correlation run as spectral
// products (one transform pair whatever the lag count), the partial
// autocorrelation follows by Durbin-Levinson recursion over the ACF,
// the algorithm that defined the PACF in the first place.

// Autocorrelate returns the sample autocorrelation of a rank-1 real
// signal at lags 0..maxLag: the mean is removed first and every value
// is normalised by the lag-zero sum (the biased estimator, the one
// whose Durbin-Levinson and Bartlett-band theory is written for).
// maxLag must lie in [0, n−1]; the result has maxLag+1 entries and
// starts at 1, unless the signal is constant: its zero power makes
// every lag 0, lag zero included.
func Autocorrelate(x *core.Array, maxLag int) (*core.Array, error) {
	const name = "Autocorrelate"
	if x.NDim() != 1 {
		return nil, base.Errf("%s: needs a rank-1 signal, got shape %s", name, base.ShapeText(x.Shape()))
	}
	if x.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex signals are not supported", name)
	}
	n := x.Len()
	if n == 0 {
		return nil, base.Errf("%s: an empty signal has no correlation", name)
	}
	if maxLag < 0 || maxLag >= n {
		return nil, base.Errf("%s: maxLag must lie in [0, %d], got %d", name, n-1, maxLag)
	}
	// Demean. The payload walk reads the same values FloatAt would
	// (widenFloats widens exactly), in the same order, so the mean
	// and every padded sample keep their bits.
	xF := x.RawFloats()
	if xF == nil || x.Strided() {
		xF = widenFloats(x)
	}
	// A non-finite sample would drive the normaliser NaN and publish
	// NaN correlations with no error, so it is refused up front (the
	// guard SolvePoissonPeriodic applies to its source).
	for i := range n {
		if v := xF[i]; math.IsNaN(v) || math.IsInf(v, 0) {
			return nil, base.Errf("%s: x holds the non-finite value %g at %d", name, v, i)
		}
	}
	mean := 0.0
	for i := range n {
		mean += xF[i]
	}
	mean /= float64(n)
	// The correlation runs as one transform pair over the zero-padded
	// signal, in place on a private buffer: the values that reach the
	// butterflies and the arithmetic that follows are the FFT/IFFT
	// path's unchanged, minus the copies into and out of arrays that
	// nobody else ever read.
	padded := make([]complex128, 2*n)
	for i := range n {
		padded[i] = complex(xF[i]-mean, 0)
	}
	transform(padded, -1)
	// |F|²: multiply by the conjugate, then back. The zero-padded
	// transform makes the circular correlation linear.
	for i := range padded {
		padded[i] = padded[i] * conj128(padded[i])
	}
	transform(padded, +1)
	// The inverse scaling, applied with IFFT's own operation order.
	m := float64(2 * n)
	for i := range padded {
		padded[i] /= complex(m, 0)
	}
	den := real(padded[0])
	out := make([]float64, maxLag+1)
	for k := range maxLag + 1 {
		if den == 0 {
			out[k] = 0
			continue
		}
		out[k] = real(padded[k]) / den
	}
	return floatsFromArrayMust(out, []int{maxLag + 1}), nil
}

// conj128 conjugates one complex number (the local spelling keeps the
// dependency surface to the packages already imported).
func conj128(z complex128) complex128 { return complex(real(z), -imag(z)) }

// CrossCorrelate returns the raw cross-correlation of two equal-length
// real signals at every lag: the result has 2n−1 entries, index i
// carrying lag i−(n−1), value r(lag) = Σ_t x[t+lag]·y[t] over the
// overlapping t. No mean removal, no normalisation: correlation as
// the inner product of shifted copies, the definition convolution and
// matched filtering share.
func CrossCorrelate(x, y *core.Array) (*core.Array, error) {
	const name = "CrossCorrelate"
	if x.NDim() != 1 || y.NDim() != 1 {
		return nil, base.Errf("%s: needs two rank-1 signals, got %s and %s",
			name, base.ShapeText(x.Shape()), base.ShapeText(y.Shape()))
	}
	if x.Dtype() == core.Complex || y.Dtype() == core.Complex {
		return nil, base.Errf("%s: complex signals are not supported", name)
	}
	n := x.Len()
	if n == 0 || y.Len() != n {
		return nil, base.Errf("%s: the signals must be equal-length non-empty, got %d and %d",
			name, n, y.Len())
	}
	// Zero-pad both signals through the raw payloads where possible;
	// the widened reads equal the FloatAt values bit for bit.
	xF := x.RawFloats()
	if xF == nil || x.Strided() {
		xF = widenFloats(x)
	}
	yF := y.RawFloats()
	if yF == nil || y.Strided() {
		yF = widenFloats(y)
	}
	// A non-finite sample would publish NaN correlations with no
	// error, so both signals are refused up front.
	for i := range n {
		if v := xF[i]; math.IsNaN(v) || math.IsInf(v, 0) {
			return nil, base.Errf("%s: x holds the non-finite value %g at %d", name, v, i)
		}
		if v := yF[i]; math.IsNaN(v) || math.IsInf(v, 0) {
			return nil, base.Errf("%s: y holds the non-finite value %g at %d", name, v, i)
		}
	}
	// The spectra multiply in place on private buffers: the values
	// reaching the transforms and the per-element products are the
	// FFT/mulSpectraConj/IFFT path's unchanged, minus the copies into
	// and out of arrays nobody else ever read.
	px := make([]complex128, 2*n)
	py := make([]complex128, 2*n)
	for i := range n {
		px[i] = complex(xF[i], 0)
		py[i] = complex(yF[i], 0)
	}
	transform(px, -1)
	transform(py, -1)
	// Fx·conj(Fy) transforms back to Σ_t x[t+lag]·y[t], the documented
	// orientation; the index mapping lands the positive lags in the
	// first half.
	// The product runs in mulSpectraConj, the array-based helper the
	// FFT path always used: the compiler's FMA fusion of the complex
	// multiply depends on the compiled shape of its surroundings, and
	// the helper's shape is the one whose product bits this result is
	// pinned to. The wrappers take the buffers without copying; the
	// product array is the one allocation of the old path this keeps.
	prod, perr := mulSpectraConj(complexFromArrayMust(px, []int{2 * n}),
		complexFromArrayMust(py, []int{2 * n}))
	if perr != nil {
		return nil, base.Errf("%s: %w", name, perr)
	}
	pc := prod.RawComplexes()
	transform(pc, +1)
	m := float64(2 * n)
	for i := range pc {
		pc[i] /= complex(m, 0)
	}
	// Lag 0 sits at index 0 of the transformed product, lag k at k and
	// negative lag −k at 2n−k; the output centres them on index n−1.
	out := make([]float64, 2*n-1)
	for k := range 2*n - 1 {
		lag := k - (n - 1)
		if lag >= 0 {
			out[k] = real(pc[lag])
		} else {
			out[k] = real(pc[2*n+lag])
		}
	}
	return floatsFromArrayMust(out, []int{2*n - 1}), nil
}

// mulSpectraConj multiplies the first spectrum by the conjugate of the
// second, element-wise. Both spectra come from the FFT and stream
// their complex payloads directly; the accessor loop stays as the
// fallback for anything else. The function is kept byte-for-byte as
// the FFT path always defined it: its compiled FMA pattern is what the
// correlation results are pinned to.
func mulSpectraConj(a, b *core.Array) (*core.Array, error) {
	as := a.RawComplexes()
	bs := b.RawComplexes()
	if as != nil && bs != nil && !a.Strided() && !b.Strided() {
		out := make([]complex128, len(as))
		for i := range out {
			out[i] = as[i] * conj128(bs[i])
		}
		return core.ComplexFromArray(out, a.Shape()...)
	}
	out := make([]complex128, a.Len())
	for i := range out {
		out[i] = a.ComplexAt(i) * conj128(b.ComplexAt(i))
	}
	return core.ComplexFromArray(out, a.Shape()...)
}

// PartialAutocorrelate returns the partial autocorrelation of a rank-1
// real signal at lags 1..maxLag, by the Durbin-Levinson recursion over
// the biased autocorrelation: pacf[k] is the last coefficient of the
// order-k AR fit, the "correlation with the intermediate lags
// regressed out" that box-jenkins identification reads. maxLag must
// lie in [1, n−2] for the recursion to stay meaningful.
func PartialAutocorrelate(x *core.Array, maxLag int) (*core.Array, error) {
	const name = "PartialAutocorrelate"
	n := x.Len()
	if maxLag < 1 || maxLag > n-2 {
		return nil, base.Errf("%s: maxLag must lie in [1, %d], got %d", name, n-2, maxLag)
	}
	acf, err := Autocorrelate(x, maxLag)
	if err != nil {
		return nil, err
	}
	// r[k] = acf at lag k (1-based recursion indexing).
	r := make([]float64, maxLag+1)
	r[0] = 1
	for k := range maxLag {
		r[k+1] = acf.FloatAt(k + 1)
	}
	phi := make([]float64, maxLag+1) // current order's coefficients
	phiPrev := make([]float64, maxLag+1)
	out := make([]float64, maxLag)
	for k := 1; k <= maxLag; k++ {
		num := r[k]
		den := 1.0
		for j := 1; j < k; j++ {
			num -= phiPrev[j] * r[k-j]
			den -= phiPrev[j] * r[j]
		}
		if den == 0 {
			return nil, base.Errf("%s: the recursion broke down at lag %d", name, k)
		}
		phi[k] = num / den
		for j := 1; j < k; j++ {
			phi[j] = phiPrev[j] - phi[k]*phiPrev[k-j]
		}
		copy(phiPrev, phi)
		out[k-1] = phi[k]
	}
	return floatsFromArrayMust(out, []int{maxLag}), nil
}