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

// daubechiesFamilies lists every family the machinery carries, with
// its vanishing-moment count.
var daubechiesFamilies = []struct {
	family  Daubechies
	moments int
	coeffs  []float64
}{
	{DB2, 2, db2Coeffs},
	{DB3, 3, db3Coeffs},
	{DB4, 4, db4Coeffs},
	{DB5, 5, db5Coeffs},
	{DB6, 6, db6Coeffs},
	{DB7, 7, db7Coeffs},
	{DB8, 8, db8Coeffs},
}

// TestDaubechiesPerfectReconstruction pins the defining property of
// the orthonormal bank: inverse(forward(x)) == x to machine
// precision, for every family, at every level count the length
// allows.
func TestDaubechiesPerfectReconstruction(t *testing.T) {
	for _, f := range daubechiesFamilies {
		L := 2 * f.moments
		for _, mult := range []int{4, 8} {
			n := L * mult
			g := core.NewGenerator(int64(f.moments)*100 + int64(mult))
			vals := make([]float64, n)
			for i := range vals {
				vals[i] = g.NormalUnit()
			}
			x := mustFloats(t, vals)
			for _, levels := range []int{1, 2, 3} {
				c, err := DaubechiesDWT(x, f.family, levels, DWTPeriodic)
				if err != nil {
					t.Fatalf("%s n=%d levels=%d: %v", f.family, n, levels, err)
				}
				back, err := DaubechiesIDWT(c, f.family, levels, DWTPeriodic)
				if err != nil {
					t.Fatalf("%s n=%d levels=%d inverse: %v", f.family, n, levels, err)
				}
				for i := range vals {
					if math.Abs(back.FloatAt(i)-vals[i]) > 1e-12 {
						t.Fatalf("%s n=%d levels=%d: reconstruction off at %d: %v vs %v",
							f.family, n, levels, i, back.FloatAt(i), vals[i])
					}
				}
			}
		}
	}
}

// TestDaubechiesEnergyPreservation pins Parseval: the periodic
// transform moves no energy between the signal and its coefficients.
func TestDaubechiesEnergyPreservation(t *testing.T) {
	for _, f := range []struct {
		family  Daubechies
		moments int
	}{
		{DB4, 4},
		{DB7, 7},
	} {
		L := 2 * f.moments
		n := 8 * L
		g := core.NewGenerator(int64(f.moments) + 50)
		vals := make([]float64, n)
		for i := range vals {
			vals[i] = g.NormalUnit()
		}
		x := mustFloats(t, vals)
		for _, levels := range []int{1, 2, 4} {
			c, err := DaubechiesDWT(x, f.family, levels, DWTPeriodic)
			if err != nil {
				t.Fatalf("%s levels=%d: %v", f.family, levels, err)
			}
			eIn, eC := 0.0, 0.0
			for i := range n {
				eIn += vals[i] * vals[i]
				eC += c.FloatAt(i) * c.FloatAt(i)
			}
			if math.Abs(eIn-eC) > 1e-11*eIn {
				t.Fatalf("%s levels=%d: energy %g vs %g", f.family, levels, eC, eIn)
			}
		}
	}
}

// TestDaubechiesFilterTables certifies every coefficient table
// against the defining conditions of the family: the √2 sum, the
// unit norm, the shift-2 orthogonality, the vanishing moments, and
// the spectral identity against Daubechies' polynomial P, so a
// mistyped digit cannot survive.
func TestDaubechiesFilterTables(t *testing.T) {
	for _, f := range daubechiesFamilies {
		h := f.coeffs
		L := len(h)
		sum, norm := 0.0, 0.0
		for _, v := range h {
			sum += v
			norm += v * v
		}
		if math.Abs(sum-math.Sqrt2) > 1e-14 {
			t.Fatalf("%s: sum %v, want √2", f.family, sum)
		}
		if math.Abs(norm-1) > 1e-14 {
			t.Fatalf("%s: norm %v, want 1", f.family, norm)
		}
		for l := 1; 2*l < L; l++ {
			s := 0.0
			for k := 0; k+2*l < L; k++ {
				s += h[k] * h[k+2*l]
			}
			if math.Abs(s) > 1e-13 {
				t.Fatalf("%s: shift-2 orthogonality at l=%d: %v", f.family, l, s)
			}
		}
		// The vanishing moments, judged relatively: the raw sums of
		// k^j-weighted terms reach magnitudes where the arithmetic
		// noise floor alone is around 1e-9.
		for j := 0; j < f.moments; j++ {
			var s, scale float64
			for k, v := range h {
				s += math.Pow(-1, float64(k)) * math.Pow(float64(k), float64(j)) * v
				scale += math.Pow(float64(k), float64(j)) * math.Abs(v)
			}
			if math.Abs(s) > 1e-12*scale {
				t.Fatalf("%s: vanishing moment %d: %v (scale %v)", f.family, j, s, scale)
			}
		}
		// The spectral identity |H(ω)|² = 2·cos^{2N}(ω/2)·P(sin²(ω/2))
		// with P(y) = Σ C(N−1+k, k)·y^k.
		P := make([]float64, f.moments)
		for k := range P {
			P[k] = binomialCoeffs(f.moments - 1 + k)[k]
		}
		for i := 0; i <= 512; i++ {
			w := math.Pi * float64(i) / 512
			var hr, hi float64
			for k, v := range h {
				hr += v * math.Cos(float64(k)*w)
				hi -= v * math.Sin(float64(k)*w)
			}
			s := math.Sin(w / 2)
			y := s * s
			pv := 0.0
			pow := 1.0
			for k := range P {
				pv += P[k] * pow
				pow *= y
			}
			c := math.Cos(w / 2)
			want := 2 * math.Pow(c*c, float64(f.moments)) * pv
			if got := hr*hr + hi*hi; math.Abs(got-want) > 1e-12 {
				t.Fatalf("%s: spectral identity at ω %d/512: %v vs %v", f.family, i, got, want)
			}
		}
	}
}

// TestDaubechiesKnownTwoLevel pins a whole 2-level decomposition,
// db2 over [1..8], against independently computed reference values
// (an explicit analysis matrix built row by row), layout included.
func TestDaubechiesKnownTwoLevel(t *testing.T) {
	x := mustFloats(t, []float64{1, 2, 3, 4, 5, 6, 7, 8})
	c, err := DaubechiesDWT(x, DB2, 2, DWTPeriodic)
	if err != nil {
		t.Fatalf("DaubechiesDWT: %v", err)
	}
	want := []float64{
		5.9019237886466849, 12.098076211353316,
		0.36602540378444015, -3.8301270189221936,
		0, 0, 0, -2.8284271247461907,
	}
	for i, w := range want {
		got := c.FloatAt(i)
		if math.Abs(got-w) > 1e-12 {
			t.Fatalf("coefficient %d: %.15g, want %.15g", i, got, w)
		}
	}
}

// TestDaubechiesConstantDetails pins the vanishing moments end to
// end: the detail bands of a constant signal are zero at every level
// and the approximation scales by 2^(levels/2), the Σh = √2 gain.
func TestDaubechiesConstantDetails(t *testing.T) {
	for _, f := range []struct {
		family  Daubechies
		moments int
	}{
		{DB2, 2},
		{DB5, 5},
		{DB8, 8},
	} {
		L := 2 * f.moments
		n := 4 * L
		for _, levels := range []int{1, 2} {
			vals := make([]float64, n)
			for i := range vals {
				vals[i] = 3.5
			}
			c, err := DaubechiesDWT(mustFloats(t, vals), f.family, levels, DWTPeriodic)
			if err != nil {
				t.Fatalf("%s levels=%d: %v", f.family, levels, err)
			}
			block := n >> levels
			wantA := 3.5 * math.Pow(math.Sqrt2, float64(levels))
			for i := range block {
				if got := c.FloatAt(i); math.Abs(got-wantA) > 1e-12 {
					t.Fatalf("%s levels=%d: approximation %d: %v, want %v",
						f.family, levels, i, got, wantA)
				}
			}
			for i := block; i < n; i++ {
				if got := c.FloatAt(i); math.Abs(got) > 1e-12 {
					t.Fatalf("%s levels=%d: detail coefficient %d = %v, want 0",
						f.family, levels, i, got)
				}
			}
		}
	}
}

// TestDaubechiesZeroPad pins the padding mode: a length no level
// tree would accept transforms on the next valid padded length, the
// padded signal's energy is conserved, and the inverse returns the
// padded length whose head is the original signal.
func TestDaubechiesZeroPad(t *testing.T) {
	g := core.NewGenerator(77)
	const n = 13
	vals := make([]float64, n)
	var eIn float64
	for i := range vals {
		vals[i] = g.NormalUnit()
		eIn += vals[i] * vals[i]
	}
	x := mustFloats(t, vals)
	// db2 has 4 taps: two levels need a multiple of 8, so 13 pads
	// to 16.
	c, err := DaubechiesDWT(x, DB2, 2, DWTZeroPad)
	if err != nil {
		t.Fatalf("DaubechiesDWT zero-pad: %v", err)
	}
	if c.Len() != 16 {
		t.Fatalf("padded coefficient length %d, want 16", c.Len())
	}
	eC := 0.0
	for i := range 16 {
		eC += c.FloatAt(i) * c.FloatAt(i)
	}
	if math.Abs(eIn-eC) > 1e-11*eIn {
		t.Fatalf("padded energy %g vs %g", eC, eIn)
	}
	back, err := DaubechiesIDWT(c, DB2, 2, DWTZeroPad)
	if err != nil {
		t.Fatalf("DaubechiesIDWT zero-pad: %v", err)
	}
	if back.Len() != 16 {
		t.Fatalf("reconstruction length %d, want the padded 16", back.Len())
	}
	for i := range vals {
		if math.Abs(back.FloatAt(i)-vals[i]) > 1e-12 {
			t.Fatalf("reconstruction off at %d: %v vs %v", i, back.FloatAt(i), vals[i])
		}
	}
}

// TestDaubechiesErrors pins the input gates: family names, level
// counts, the length contract of both modes, and the shapes.
func TestDaubechiesErrors(t *testing.T) {
	x := mustFloats(t, []float64{1, 2, 3, 4, 5, 6, 7, 8})
	if _, err := DaubechiesDWT(x, "db1", 1, DWTPeriodic); err == nil {
		t.Error("db1 accepted; Haar keeps its own transform")
	}
	if _, err := DaubechiesDWT(x, "db9", 1, DWTPeriodic); err == nil {
		t.Error("db9 accepted")
	}
	if _, err := DaubechiesDWT(x, "", 1, DWTPeriodic); err == nil {
		t.Error("an empty family name accepted")
	}
	if _, err := DaubechiesDWT(x, DB2, 0, DWTPeriodic); err == nil {
		t.Error("zero levels accepted")
	}
	// db2 has 4 taps: 8 samples support 2 levels, not 3.
	if _, err := DaubechiesDWT(x, DB2, 3, DWTPeriodic); err == nil {
		t.Error("levels deeper than the filter length accepted")
	}
	// 12 samples hold three 4-tap blocks at level 1 but 6 are not a
	// multiple of 4 at level 2.
	twelve := mustFloats(t, make([]float64, 12))
	if _, err := DaubechiesDWT(twelve, DB2, 1, DWTPeriodic); err != nil {
		t.Errorf("12 samples at 1 level refused: %v", err)
	}
	if _, err := DaubechiesDWT(twelve, DB2, 2, DWTPeriodic); err == nil {
		t.Error("a block of 6 accepted against a 4-tap filter")
	}
	// 14 samples are not a multiple of 4 even at level 1.
	if _, err := DaubechiesDWT(mustFloats(t, make([]float64, 14)), DB2, 1, DWTPeriodic); err == nil {
		t.Error("14 samples accepted against a 4-tap filter")
	}
	// The zero-padded mode takes exactly those.
	if _, err := DaubechiesDWT(mustFloats(t, make([]float64, 14)), DB2, 1, DWTZeroPad); err != nil {
		t.Errorf("zero-pad refused a short length: %v", err)
	}
	if _, err := DaubechiesDWT(mustFloats(t, make([]float64, 14)), DB2, 61, DWTZeroPad); err == nil {
		t.Error("61 zero-pad levels accepted")
	}
	// db8's 16 taps at 60 levels would need a padded length of
	// 16·2^59, which wraps past the addressable range: refused.
	if _, err := DaubechiesDWT(mustFloats(t, make([]float64, 8)), DB8, 60, DWTZeroPad); err == nil {
		t.Error("a padded length past the addressable range accepted")
	}
	// The band between the wrapped corner and the shallow depths holds
	// spans the shift survives but no allocator does: each must refuse
	// rather than panic in the work-buffer make.
	for _, tc := range []struct {
		family Daubechies
		levels int
		why    string
	}{
		{DB2, 60, "2^61 elements"},
		{DB8, 59, "2^62 elements"},
		{DB8, 50, "2^53 elements"},
		{DB8, 43, "2^46 elements, one past the makeslice ceiling"},
	} {
		if _, err := DaubechiesDWT(mustFloats(t, make([]float64, 8)), tc.family, tc.levels, DWTZeroPad); err == nil {
			t.Errorf("%s at %d levels accepted a padded span of %s", tc.family, tc.levels, tc.why)
		}
	}
	if _, err := DaubechiesIDWT(x, DB2, 0, DWTPeriodic); err == nil {
		t.Error("zero levels accepted by the inverse")
	}
	if _, err := DaubechiesIDWT(x, DB2, 1, DWTMode(9)); err == nil {
		t.Error("unknown boundary mode accepted by the inverse")
	}
	if _, err := DaubechiesDWT(x, DB2, 1, DWTMode(9)); err == nil {
		t.Error("unknown boundary mode accepted")
	}
	if _, err := DaubechiesDWT(mustFloats(t, []float64{1, 2, 3, 4}, 2, 2), DB2, 1, DWTPeriodic); err == nil {
		t.Error("rank 2 accepted")
	}
	if _, err := DaubechiesDWT(mustComplexes(t, []complex128{1, 2, 3, 4, 5, 6, 7, 8}, 8), DB2, 1, DWTPeriodic); err == nil {
		t.Error("complex accepted")
	}
	if _, err := DaubechiesDWT(mustFloats(t, nil), DB2, 1, DWTPeriodic); err == nil {
		t.Error("empty accepted")
	}
	// The inverse repeats the length contract on its own input.
	c, err := DaubechiesDWT(x, DB2, 1, DWTPeriodic)
	if err != nil {
		t.Fatal(err)
	}
	bad := mustFloats(t, make([]float64, 10))
	if _, err := DaubechiesIDWT(bad, DB2, 1, DWTPeriodic); err == nil {
		t.Error("a length-10 coefficient vector accepted")
	}
	if _, err := DaubechiesIDWT(c, "db7", 1, DWTPeriodic); err == nil {
		t.Error("inverting db2 coefficients as db7 accepted")
	}
}