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

package signal

import (
	"fmt"
	"testing"
)

// Channel-blocked kernel benchmarks at inference-like sizes. The
// channel count is the axis the kernels divide work along, so each
// size climbs a cache level: 4 channels stay in L1, 32 cross L2 and
// 128 make the whole-volume walk pay for memory bandwidth.

// BenchmarkConv3D measures a (1, C, 64, 64, 64) input through C×C
// kernels of 3×3×3, stride 1, padding 1.
func BenchmarkConv3D(b *testing.B) {
	for _, ch := range []int{4, 32, 128} {
		b.Run(fmt.Sprintf("%d channels", ch), func(b *testing.B) {
			input := benchSig(b, 1, ch*64*64*64, 1, ch, 64, 64, 64)
			kernel := benchSig(b, 2, ch*ch*27, ch, ch, 3, 3, 3)
			b.ReportAllocs()
			for b.Loop() {
				if _, err := Conv3D(input, kernel, nil, 1, [3]int{1, 1, 1}, [3]int{1, 1, 1}); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}

// BenchmarkConvTranspose2D measures a (1, C, 128, 128) input through
// C×C kernels of 3×3 at the upsampling stride of 2, padding 1.
func BenchmarkConvTranspose2D(b *testing.B) {
	for _, ch := range []int{4, 32, 128} {
		b.Run(fmt.Sprintf("%d channels", ch), func(b *testing.B) {
			input := benchSig(b, 3, ch*128*128, 1, ch, 128, 128)
			kernel := benchSig(b, 4, ch*ch*9, ch, ch, 3, 3)
			b.ReportAllocs()
			for b.Loop() {
				if _, err := ConvTranspose2D(input, kernel, nil, 2, 1); err != nil {
					b.Fatal(err)
				}
			}
		})
	}
}