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

// TestAutocorrelateAR1 pins the ACF of an AR(1) process against its
// geometric theoretical decay.
func TestAutocorrelateAR1(t *testing.T) {
	const (
		phi = 0.6
		n   = 20000
	)
	g := core.NewGenerator(23)
	x := make([]float64, n)
	noise := 0.0
	for i := range n {
		noise = phi*noise + g.NormalUnit()
		x[i] = noise
	}
	xArr, err := core.FromFloats(x, n)
	if err != nil {
		t.Fatalf("FromFloats: %v", err)
	}
	acf, err := Autocorrelate(xArr, 5)
	if err != nil {
		t.Fatalf("Autocorrelate: %v", err)
	}
	if math.Abs(acf.FloatAt(0)-1) > 1e-12 {
		t.Fatalf("acf[0] = %.12f, want 1", acf.FloatAt(0))
	}
	for k := 1; k <= 5; k++ {
		want := math.Pow(phi, float64(k))
		if math.Abs(acf.FloatAt(k)-want) > 0.06 {
			t.Fatalf("acf[%d] = %.4f, want %.4f", k, acf.FloatAt(k), want)
		}
	}
}

// TestAutocorrelateBrute pins the ACF against the direct sum.
func TestAutocorrelateBrute(t *testing.T) {
	x := []float64{3, -1, 4, -1.5, 2, -0.5, 1, 2}
	n := len(x)
	xArr, _ := core.FromFloats(x, n)
	acf, err := Autocorrelate(xArr, n-1)
	if err != nil {
		t.Fatalf("Autocorrelate: %v", err)
	}
	mean := 0.0
	for _, v := range x {
		mean += v
	}
	mean /= float64(n)
	num := 0.0
	for _, v := range x {
		num += (v - mean) * (v - mean)
	}
	for k := range n {
		s := 0.0
		for t := 0; t+k < n; t++ {
			s += (x[t] - mean) * (x[t+k] - mean)
		}
		want := s / num
		if math.Abs(acf.FloatAt(k)-want) > 1e-12 {
			t.Fatalf("acf[%d] = %.12f, direct sum %.12f", k, acf.FloatAt(k), want)
		}
	}
}

// TestCrossCorrelateBrute pins the XCF against the direct sum on the
// documented lag convention.
func TestCrossCorrelateBrute(t *testing.T) {
	x := []float64{2, -1, 3, 0.5, -2}
	y := []float64{1, 4, -2, 0.5, 3}
	n := len(x)
	xArr, _ := core.FromFloats(x, n)
	yArr, _ := core.FromFloats(y, n)
	xcf, err := CrossCorrelate(xArr, yArr)
	if err != nil {
		t.Fatalf("CrossCorrelate: %v", err)
	}
	if xcf.Len() != 2*n-1 {
		t.Fatalf("length %d, want %d", xcf.Len(), 2*n-1)
	}
	for k := range 2*n - 1 {
		lag := k - (n - 1)
		s := 0.0
		for t := 0; t+lag < n && t < n; t++ {
			if t+lag < 0 {
				continue
			}
			s += x[t+lag] * y[t]
		}
		if math.Abs(xcf.FloatAt(k)-s) > 1e-12 {
			t.Fatalf("xcf[lag %d] = %.12f, direct sum %.12f", lag, xcf.FloatAt(k), s)
		}
	}
}

// TestPartialAutocorrelateAR1 pins the PACF signature of an AR(1):
// lag 1 estimates phi, the rest are noise around zero.
func TestPartialAutocorrelateAR1(t *testing.T) {
	const (
		phi = 0.7
		n   = 30000
	)
	g := core.NewGenerator(31)
	x := make([]float64, n)
	e := 0.0
	for i := range n {
		e = phi*e + g.NormalUnit()
		x[i] = e
	}
	xArr, _ := core.FromFloats(x, n)
	pacf, err := PartialAutocorrelate(xArr, 6)
	if err != nil {
		t.Fatalf("PartialAutocorrelate: %v", err)
	}
	if math.Abs(pacf.FloatAt(0)-phi) > 0.05 {
		t.Fatalf("pacf[1] = %.4f, want %.2f", pacf.FloatAt(0), phi)
	}
	for k := 2; k <= 6; k++ {
		if math.Abs(pacf.FloatAt(k-1)) > 0.05 {
			t.Fatalf("pacf[%d] = %.4f, an AR(1) cuts off after lag 1", k, pacf.FloatAt(k-1))
		}
	}
}

// TestCorrelateErrors pins the input gates.
func TestCorrelateErrors(t *testing.T) {
	x, _ := core.FromFloats([]float64{1, 2, 3}, 3)
	if _, err := Autocorrelate(x, 3); err == nil {
		t.Error("maxLag ≥ n accepted")
	}
	m, _ := core.FromFloats([]float64{1, 2}, 1, 2)
	if _, err := Autocorrelate(m, 0); err == nil {
		t.Error("rank-2 signal accepted")
	}
	y, _ := core.FromFloats([]float64{1, 2}, 2)
	if _, err := CrossCorrelate(x, y); err == nil {
		t.Error("length mismatch accepted")
	}
	if _, err := PartialAutocorrelate(x, 5); err == nil {
		t.Error("oversized PACF maxLag accepted")
	}
}