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

package signal

import (
	"math"
	"math/rand/v2"
	"testing"

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

func TestConv2DBadBias(t *testing.T) {
	// Conv2D must reject a bias whose length does not match the output
	// channels, like the other convolution entry points.
	in := mustFromFloats(t, []float64{1, 2, 3, 4}, 1, 1, 2, 2)
	ker := mustFromFloats(t, []float64{1, 0, 0, 1}, 1, 1, 2, 2)
	bad := mustFromFloats(t, []float64{1, 2, 3}, 3) // one output channel, 3 biases
	if _, err := Conv2D(in, ker, bad, 1, 0); err == nil {
		t.Error("Conv2D: expected error for mismatched bias")
	}
	good := mustFromFloats(t, []float64{1}, 1)
	if _, err := Conv2D(in, ker, good, 1, 0); err != nil {
		t.Fatalf("Conv2D with valid bias: %v", err)
	}
}

func TestConv1D(t *testing.T) {
	// 1×1×5 input, 1×1×3 kernel.
	in := mustFromFloats(t, []float64{1, 2, 3, 4, 5}, 1, 1, 5)
	ker := mustFromFloats(t, []float64{1, 0, -1}, 1, 1, 3)
	got, err := Conv1D(in, ker, nil, 1, 0, 1)
	if err != nil {
		t.Fatal(err)
	}
	if got.Shape()[2] != 3 {
		t.Errorf("Conv1D output L: %d, want 3", got.Shape()[2])
	}
	// (1,2,3) convolved with (1,0,-1) = [1*1 + 2*0 + 3*(-1), 2*1 + 3*0 + 4*(-1), 3*1 + 4*0 + 5*(-1)]
	// = [-2, -2, -2]
	for i, w := range []float64{-2, -2, -2} {
		v, _ := core.FloatAt(got, 0, 0, i)
		if v != w {
			t.Errorf("Conv1D [%d]: got %v, want %v", i, v, w)
		}
	}
}

func TestConv3D(t *testing.T) {
	// 1×1×1×2×2 input, 1×1×1×2×2 kernel (basic case).
	in := mustFromFloats(t, []float64{1, 2, 3, 4}, 1, 1, 1, 2, 2)
	ker := mustFromFloats(t, []float64{1, 1, 1, 1}, 1, 1, 1, 2, 2)
	got, err := Conv3D(in, ker, nil, 1, [3]int{0, 0, 0}, [3]int{1, 1, 1})
	if err != nil {
		t.Fatal(err)
	}
	if got.Shape()[0] != 1 || got.Shape()[1] != 1 || got.Shape()[2] != 1 || got.Shape()[3] != 1 || got.Shape()[4] != 1 {
		t.Errorf("Conv3D shape: %v", got.Shape())
	}
	// Sum of 1+2+3+4 = 10.
	v, _ := core.FloatAt(got, 0, 0, 0, 0, 0)
	if v != 10 {
		t.Errorf("Conv3D: got %v, want 10", v)
	}
}

func TestConv2DGroups(t *testing.T) {
	// 2 channels, groups=2 gives depthwise convolution.
	in := mustFromFloats(t, []float64{
		1, 2, 3, 4,
		5, 6, 7, 8,
	}, 1, 2, 2, 2)
	ker := mustFromFloats(t, []float64{
		1, 0, // kernel for channel 0
		0, 1, // kernel for channel 1
	}, 2, 1, 1, 2)
	got, err := Conv2DGroups(in, ker, nil, 1, 0, 2)
	if err != nil {
		t.Fatal(err)
	}
	if got.Shape()[1] != 2 {
		t.Errorf("Conv2DGroups C_out: %d, want 2", got.Shape()[1])
	}
	// Channel 0: kernel [[1,0],[0,1]] applied to [[1,2],[3,4]] = [[1+4, 2+0],[3+0, 4+0]] = [[5,2],[3,4]]
	// Wait: Conv2D with kH=1,kW=2 means 1-row × 2-col kernel. Output H=2, W=1.
	// Channel 0 input [[1,2],[3,4]] conv with [[1,0]] (kH=1, kW=2) gives [[1·1+2·0, ...], [3·1+4·0, ...]] = [[1, ?], [3, ?]].
	v00, _ := core.FloatAt(got, 0, 0, 0, 0)
	v10, _ := core.FloatAt(got, 0, 0, 1, 0)
	// Channel 1: kernel [[0,1]] applied to [[5,6],[7,8]] = [[5·0+6·1, ...], [7·0+8·1, ...]] = [[6, ?], [8, ?]].
	v01, _ := core.FloatAt(got, 0, 1, 0, 0)
	v11, _ := core.FloatAt(got, 0, 1, 1, 0)
	if v00 != 1 || v10 != 3 || v01 != 6 || v11 != 8 {
		t.Errorf("Conv2DGroups: got ch0 [[%v,?],[%v,?]], ch1 [[%v,?],[%v,?]], want [[1,_],[3,_]], [[6,_],[8,_]]",
			v00, v10, v01, v11)
	}
}

func TestConvTranspose2D(t *testing.T) {
	// 1×1×1×1 input, 1×1×2×2 kernel: output should be 2×2.
	in := mustFromFloats(t, []float64{1}, 1, 1, 1, 1)
	ker := mustFromFloats(t, []float64{
		1, 2,
		3, 4,
	}, 1, 1, 2, 2)
	got, err := ConvTranspose2D(in, ker, nil, 1, 0)
	if err != nil {
		t.Fatal(err)
	}
	if got.Shape()[2] != 2 || got.Shape()[3] != 2 {
		t.Errorf("ConvTranspose2D shape: %v", got.Shape())
	}
	// With single input and identity-position kernel, output == kernel.
	for i, w := range []float64{1, 2, 3, 4} {
		v, _ := core.FloatAt(got, 0, 0, i/2, i%2)
		if math.Abs(v-w) > 1e-9 {
			t.Errorf("ConvTranspose2D [%d]: got %v, want %v", i, v, w)
		}
	}
}

func TestMaxPool1D(t *testing.T) {
	in := mustFromFloats(t, []float64{1, 3, 2, 4, 1}, 1, 1, 5)
	got, err := MaxPool1D(in, 2, 2, 0)
	if err != nil {
		t.Fatal(err)
	}
	// 2 windows: [1,3] gives 3, [2,4] gives 4.
	v0, _ := core.FloatAt(got, 0, 0, 0)
	v1, _ := core.FloatAt(got, 0, 0, 1)
	if v0 != 3 || v1 != 4 {
		t.Errorf("MaxPool1D: got %v %v, want 3 4", v0, v1)
	}
}

func TestAvgPool3D(t *testing.T) {
	in := mustFromFloats(t, []float64{1, 2, 3, 4, 5, 6, 7, 8}, 1, 1, 2, 2, 2)
	got, err := AvgPool3D(in, [3]int{2, 2, 2}, [3]int{2, 2, 2}, [3]int{0, 0, 0}, true)
	if err != nil {
		t.Fatal(err)
	}
	if got.Shape()[2] != 1 || got.Shape()[3] != 1 || got.Shape()[4] != 1 {
		t.Errorf("AvgPool3D shape: %v", got.Shape())
	}
	v, _ := core.FloatAt(got, 0, 0, 0, 0, 0)
	// Average of all 8 = 36 / 8 = 4.5.
	if math.Abs(v-4.5) > 1e-9 {
		t.Errorf("AvgPool3D: got %v, want 4.5", v)
	}
}

func TestAdaptiveMaxPool2D(t *testing.T) {
	in := mustFromFloats(t, []float64{
		1, 2, 3, 4,
		5, 6, 7, 8,
		9, 10, 11, 12,
		13, 14, 15, 16,
	}, 1, 1, 4, 4)
	got, err := AdaptiveMaxPool2D(in, 2, 2)
	if err != nil {
		t.Fatal(err)
	}
	v00, _ := core.FloatAt(got, 0, 0, 0, 0)
	v01, _ := core.FloatAt(got, 0, 0, 0, 1)
	v11, _ := core.FloatAt(got, 0, 0, 1, 1)
	if v00 != 6 || v01 != 8 || v11 != 16 {
		t.Errorf("AdaptiveMaxPool2D: got %v %v %v, want 6 8 16", v00, v01, v11)
	}
}

func TestGlobalAvgPool2D(t *testing.T) {
	in := mustFromFloats(t, []float64{
		1, 2, 3, 4,
		5, 6, 7, 8,
	}, 1, 1, 2, 4)
	got, err := GlobalAvgPool2D(in)
	if err != nil {
		t.Fatal(err)
	}
	v, _ := core.FloatAt(got, 0, 0, 0, 0)
	// Average of 1..8 = 36/8 = 4.5.
	if math.Abs(v-4.5) > 1e-9 {
		t.Errorf("GlobalAvgPool2D: got %v, want 4.5", v)
	}
}

func TestGlobalMaxPool2D(t *testing.T) {
	in := mustFromFloats(t, []float64{
		1, 2, 3, 4,
		5, 6, 7, 8,
	}, 1, 1, 2, 4)
	got, err := GlobalMaxPool2D(in)
	if err != nil {
		t.Fatal(err)
	}
	v, _ := core.FloatAt(got, 0, 0, 0, 0)
	if v != 8 {
		t.Errorf("GlobalMaxPool2D: got %v, want 8", v)
	}
}

func TestGlobalAvgPool1D(t *testing.T) {
	in := mustFromFloats(t, []float64{1, 2, 3, 4, 5, 6}, 1, 1, 6)
	got, err := GlobalAvgPool1D(in)
	if err != nil {
		t.Fatal(err)
	}
	if got.Shape()[2] != 1 {
		t.Errorf("GlobalAvgPool1D L_out: %d", got.Shape()[2])
	}
	v, _ := core.FloatAt(got, 0, 0, 0)
	if math.Abs(v-3.5) > 1e-9 {
		t.Errorf("GlobalAvgPool1D: got %v, want 3.5", v)
	}
}

func TestGlobalMaxPool1D(t *testing.T) {
	in := mustFromFloats(t, []float64{1, 5, 3, 4, 2, 6}, 1, 1, 6)
	got, err := GlobalMaxPool1D(in)
	if err != nil {
		t.Fatal(err)
	}
	v, _ := core.FloatAt(got, 0, 0, 0)
	if v != 6 {
		t.Errorf("GlobalMaxPool1D: got %v, want 6", v)
	}
}

func TestGlobalAvgPool3D(t *testing.T) {
	in := mustFromFloats(t, []float64{
		1, 2, 3, 4, 5, 6, 7, 8,
	}, 1, 1, 2, 2, 2)
	got, err := GlobalAvgPool3D(in)
	if err != nil {
		t.Fatal(err)
	}
	v, _ := core.FloatAt(got, 0, 0, 0, 0, 0)
	// Average of 1..8 = 4.5.
	if math.Abs(v-4.5) > 1e-9 {
		t.Errorf("GlobalAvgPool3D: got %v, want 4.5", v)
	}
}

func TestGlobalMaxPool3D(t *testing.T) {
	in := mustFromFloats(t, []float64{
		1, 2, 3, 4, 5, 6, 7, 8,
	}, 1, 1, 2, 2, 2)
	got, err := GlobalMaxPool3D(in)
	if err != nil {
		t.Fatal(err)
	}
	v, _ := core.FloatAt(got, 0, 0, 0, 0, 0)
	if v != 8 {
		t.Errorf("GlobalMaxPool3D: got %v, want 8", v)
	}
}

func TestAdaptiveMaxPool1D(t *testing.T) {
	in := mustFromFloats(t, []float64{1, 2, 3, 4, 5, 6}, 1, 1, 6)
	got, err := AdaptiveMaxPool1D(in, 2)
	if err != nil {
		t.Fatal(err)
	}
	if got.Shape()[2] != 2 {
		t.Errorf("AdaptiveMaxPool1D L_out: %d", got.Shape()[2])
	}
	// Window [0..3) gives max(1,2,3) = 3; [3..6) gives max(4,5,6) = 6.
	v0, _ := core.FloatAt(got, 0, 0, 0)
	v1, _ := core.FloatAt(got, 0, 0, 1)
	if v0 != 3 || v1 != 6 {
		t.Errorf("AdaptiveMaxPool1D: got %v %v, want 3 6", v0, v1)
	}
}

func TestAdaptiveAvgPool1D(t *testing.T) {
	in := mustFromFloats(t, []float64{1, 2, 3, 4, 5, 6}, 1, 1, 6)
	got, err := AdaptiveAvgPool1D(in, 2)
	if err != nil {
		t.Fatal(err)
	}
	// [1..4) gives 6/3 = 2; [4..7) gives 15/3 = 5.
	v0, _ := core.FloatAt(got, 0, 0, 0)
	v1, _ := core.FloatAt(got, 0, 0, 1)
	if math.Abs(v0-2) > 1e-9 || math.Abs(v1-5) > 1e-9 {
		t.Errorf("AdaptiveAvgPool1D: got %v %v, want 2 5", v0, v1)
	}
}

// TestAdaptivePoolCeilWindows pins the floor-start, ceil-end window
// convention on sizes that do not divide: window o covers
// floor(o·in/out) to ceil((o+1)·in/out), so every input sample lands
// in some window (the old floor ends silently dropped input samples).
func TestAdaptivePoolCeilWindows(t *testing.T) {
	in := mustFromFloats(t, []float64{1, 2, 3, 4, 5}, 1, 1, 5)
	maxGot, err := AdaptiveMaxPool1D(in, 2)
	if err != nil {
		t.Fatal(err)
	}
	// Windows [0,3) and [2,5): max 3 and 5 (the floor convention gave 2).
	v0, _ := core.FloatAt(maxGot, 0, 0, 0)
	v1, _ := core.FloatAt(maxGot, 0, 0, 1)
	if v0 != 3 || v1 != 5 {
		t.Errorf("AdaptiveMaxPool1D 5 to 2: got %v %v, want 3 5", v0, v1)
	}
	avgGot, err := AdaptiveAvgPool1D(in, 2)
	if err != nil {
		t.Fatal(err)
	}
	// (1+2+3)/3 and (3+4+5)/3.
	a0, _ := core.FloatAt(avgGot, 0, 0, 0)
	a1, _ := core.FloatAt(avgGot, 0, 0, 1)
	if math.Abs(a0-2) > 1e-9 || math.Abs(a1-4) > 1e-9 {
		t.Errorf("AdaptiveAvgPool1D 5 to 2: got %v %v, want 2 4", a0, a1)
	}
	// Upsampling repeats inputs: 2 to 4 gives [0,1),[0,1),[1,2),[1,2).
	two := mustFromFloats(t, []float64{7, 9}, 1, 1, 2)
	up, err := AdaptiveMaxPool1D(two, 4)
	if err != nil {
		t.Fatal(err)
	}
	wantUp := []float64{7, 7, 9, 9}
	for i, w := range wantUp {
		if v, _ := core.FloatAt(up, 0, 0, i); v != w {
			t.Errorf("AdaptiveMaxPool1D 2 to 4 [%d]: got %v, want %v", i, v, w)
		}
	}
	// The 2-D twins follow the same ceil ends.
	sq := mustFromFloats(t, []float64{
		1, 2, 3, 4, 5,
		6, 7, 8, 9, 10,
		11, 12, 13, 14, 15,
		16, 17, 18, 19, 20,
		21, 22, 23, 24, 25,
	}, 1, 1, 5, 5)
	max2D, err := AdaptiveMaxPool2D(sq, 2, 2)
	if err != nil {
		t.Fatal(err)
	}
	// Windows are rows/cols [0,3) and [2,5): the quadrant maxima are
	// 13, 15, 23, 25.
	want2D := []float64{13, 15, 23, 25}
	for i, w := range want2D {
		oy, ox := i/2, i%2
		if v, _ := core.FloatAt(max2D, 0, 0, oy, ox); v != w {
			t.Errorf("AdaptiveMaxPool2D 5 to 2 [%d,%d]: got %v, want %v", oy, ox, v, w)
		}
	}
	avg2D, err := AdaptiveAvgPool2D(sq, 2, 2)
	if err != nil {
		t.Fatal(err)
	}
	// Top-left window average (1+2+3+6+7+8+11+12+13)/9 = 7.
	if v, _ := core.FloatAt(avg2D, 0, 0, 0, 0); math.Abs(v-7) > 1e-9 {
		t.Errorf("AdaptiveAvgPool2D 5 to 2 [0,0]: got %v, want 7", v)
	}
}

// TestPoolAndConvNumeratorGuard pins the output-size contract: a
// kernel that exceeds the padded input must be an error, not a
// 1-wide output from a negative numerator truncating toward zero.
func TestPoolAndConvNumeratorGuard(t *testing.T) {
	in4 := mustFromFloats(t, make([]float64, 16), 1, 1, 4, 4)
	if _, err := MaxPool2D(in4, 5, 2, 0); err == nil {
		t.Error("MaxPool2D: expected an error when the kernel exceeds the input")
	}
	if _, err := AvgPool2D(in4, 5, 2, 0, true); err == nil {
		t.Error("AvgPool2D: expected an error when the kernel exceeds the input")
	}
	lane := mustFromFloats(t, make([]float64, 4), 1, 1, 4)
	if _, err := MaxPool1D(lane, 5, 2, 0); err == nil {
		t.Error("MaxPool1D: expected an error when the kernel exceeds the input")
	}
	cube := mustFromFloats(t, make([]float64, 8), 1, 1, 2, 2, 2)
	if _, err := MaxPool3D(cube, [3]int{5, 1, 1}, [3]int{2, 1, 1}, [3]int{0, 0, 0}); err == nil {
		t.Error("MaxPool3D: expected an error when the kernel exceeds the input")
	}
	kernel, _ := core.FromFloats(make([]float64, 25), 1, 1, 5, 5)
	if _, err := Conv2D(in4, kernel, nil, 2, 0); err == nil {
		t.Error("Conv2D: expected an error when the kernel exceeds the input")
	}
	k1, _ := core.FromFloats(make([]float64, 5), 1, 1, 5)
	if _, err := Conv1D(lane, k1, nil, 2, 0, 1); err == nil {
		t.Error("Conv1D: expected an error when the kernel exceeds the input")
	}
	k3, _ := core.FromFloats(make([]float64, 125), 1, 1, 5, 5, 5)
	if _, err := Conv3D(cube, k3, nil, 2, [3]int{0, 0, 0}, [3]int{1, 1, 1}); err == nil {
		t.Error("Conv3D: expected an error when the kernel exceeds the input")
	}
}

// TestConv2DGroupsPartialChannelBlock pins the 2-D kernel's
// output-channel blocking when the last block of a group is partial:
// six output channels in two groups of three, on a row wide enough that
// one block covers two channels, so the group's third channel is a
// block of its own. Reading the block count as a floor instead of a
// ceiling drops that trailing block and leaves the channel at the zero
// the output array was initialised with, which the scalar walk below
// sees at once.
func TestConv2DGroupsPartialChannelBlock(t *testing.T) {
	const (
		n, cIn, hIn, wIn = 1, 2, 3, 2000
		groups           = 2
		cOut, kH, kW     = 6, 2, 3
		stride, padding  = 1, 1
	)
	cInPerGroup := cIn / groups
	cOutPerGroup := cOut / groups
	if block := convAccCap / wIn; block < 1 || block >= cOutPerGroup {
		t.Fatalf("the case needs a partial block: a block of %d against %d channels per group", block, cOutPerGroup)
	}
	rng := rand.New(rand.NewPCG(11, 13))
	in := make([]float64, n*cIn*hIn*wIn)
	for i := range in {
		in[i] = rng.NormFloat64()
	}
	ker := make([]float64, cOut*cInPerGroup*kH*kW)
	for i := range ker {
		ker[i] = rng.NormFloat64()
	}
	bias := make([]float64, cOut)
	for i := range bias {
		bias[i] = rng.NormFloat64()
	}
	x := mustFromFloats(t, in, n, cIn, hIn, wIn)
	k := mustFromFloats(t, ker, cOut, cInPerGroup, kH, kW)
	b := mustFromFloats(t, bias, cOut)
	got, err := Conv2DGroups(x, k, b, stride, padding, groups)
	if err != nil {
		t.Fatalf("Conv2DGroups: %v", err)
	}
	hOut := (hIn+2*padding-kH)/stride + 1
	wOut := (wIn+2*padding-kW)/stride + 1
	if got.Shape()[1] != cOut || got.Shape()[2] != hOut || got.Shape()[3] != wOut {
		t.Fatalf("shape %v, want [%d %d %d %d]", got.Shape(), n, cOut, hOut, wOut)
	}
	// A scalar walk in the same (input channel, tap row, tap column)
	// order the kernel accumulates its taps in.
	worst, scale := 0.0, 0.0
	for oc := range cOut {
		g := oc / cOutPerGroup
		live := false
		for oh := range hOut {
			for ow := range wOut {
				want := 0.0
				for ic := range cInPerGroup {
					for kh := range kH {
						ih := oh*stride + kh - padding
						if ih < 0 || ih >= hIn {
							continue
						}
						for kw := range kW {
							iw := ow*stride + kw - padding
							if iw < 0 || iw >= wIn {
								continue
							}
							want += in[((g*cInPerGroup+ic)*hIn+ih)*wIn+iw] *
								ker[((oc*cInPerGroup+ic)*kH+kh)*kW+kw]
						}
					}
				}
				want += bias[oc]
				gv := got.FloatAt(oc*hOut*wOut + oh*wOut + ow)
				if gv != 0 {
					live = true
				}
				if d := math.Abs(gv - want); d > worst {
					worst = d
				}
				if a := math.Abs(want); a > scale {
					scale = a
				}
			}
		}
		if !live {
			t.Errorf("output channel %d (group %d, the group's last channel) is identically zero", oc, g)
		}
	}
	if worst > 1e-12*scale {
		t.Fatalf("worst deviation from the scalar walk %.6g (scale %.6g), want the blocked sweep to agree with it", worst, scale)
	}
}