// 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"
	"sourcedock.dev/petrbalvin/tensor/internal/engine"
)

// Pooling operations for 2-D feature maps. All functions take
// NCHW input (N, C, H, W) and return (N, C, H_out, W_out). Padding is
// a single int applied to both spatial dims; stride is a single int.
// Padding must stay below the kernel: a window entirely inside the
// padding holds no data, and a max would answer −Inf there, an
// average 0.
//
// The kernels hold two constraints fixed:
//
//   - Element order. Average pooling adds its window elements in the
//     (kernel height, kernel width) nesting and divides once, with the
//     countIncludePad divisor chosen exactly as before; any other
//     order or divisor changes bits. Max pooling is order free but
//     keeps the same walk for a single shared code path, and a NaN in
//     a window wins: v > best is false for NaN, so the comparison
//     would silently drop it and answer the largest finite neighbour
//     instead.
//   - Parallel split. Work is distributed over output rows, which own
//     disjoint output elements, so no split can influence a result.
//
// Non-float inputs widen once into a scratch payload (see
// widenFloats): FloatAt widens exactly, so the values
// compared and summed are bit-identical to the per-element accessor
// reads.

// parallelWorkBudget is the number of element updates a worker should
// carry before splitting a kernel pays for the worker's spawn. Every
// scheduling floor in the package is derived from it: an item body
// costing a few hundred nanoseconds still leaves a worker below the
// budget on a tiny input, where the split would spawn goroutines to do
// less work than the spawn itself costs.
const parallelWorkBudget = 1 << 12

// workFloorFor returns the smallest number of items a worker should
// carry when one item costs about workPerItem element updates. It is
// the floor argument engine.ParallelMin wants: below it the kernel runs
// whole on the calling goroutine, and above it the split is the one
// engine.Parallel computes, so the worker choice and the chunk
// boundaries of a split kernel are unchanged.
func workFloorFor(workPerItem int) int {
	if workPerItem < 1 {
		return parallelWorkBudget
	}
	return max(1, parallelWorkBudget/workPerItem)
}

// MaxPool2D returns the maximum over each kernel-sized window. A NaN
// in a window propagates: the maximum answers NaN.
func MaxPool2D(input *core.Array, kernel, stride, padding int) (*core.Array, error) {
	return pool2D(input, kernel, stride, padding, true, false)
}

// AvgPool2D returns the average over each kernel-sized window. When
// countIncludePad is true the divisor is kH*kW; otherwise it's the
// number of non-padded elements in the window.
func AvgPool2D(input *core.Array, kernel, stride, padding int, countIncludePad bool) (*core.Array, error) {
	return pool2D(input, kernel, stride, padding, false, countIncludePad)
}

// MaxPool1D returns the maximum over each kernel-sized window of a
// 3-D tensor (N, C, L). A NaN in a window propagates: the maximum
// answers NaN.
func MaxPool1D(input *core.Array, kernel, stride, padding int) (*core.Array, error) {
	return pool1D(input, kernel, stride, padding, true, false)
}

// AvgPool1D returns the average over each kernel-sized window of a
// 3-D tensor (N, C, L). When countIncludePad is true the divisor is
// the full kernel size; otherwise it is the number of non-padded
// elements in the window.
func AvgPool1D(input *core.Array, kernel, stride, padding int, countIncludePad bool) (*core.Array, error) {
	return pool1D(input, kernel, stride, padding, false, countIncludePad)
}

// MaxPool3D returns the maximum over each kernel-sized window of a
// 5-D tensor (N, C, D, H, W). A NaN in a window propagates: the
// maximum answers NaN.
func MaxPool3D(input *core.Array, kernel [3]int, stride [3]int, padding [3]int) (*core.Array, error) {
	return pool3D(input, kernel, stride, padding, true, false)
}

// AvgPool3D returns the average over each kernel-sized window of a
// 5-D tensor (N, C, D, H, W). When countIncludePad is true the
// divisor is the full kernel size; otherwise it is the number of
// non-padded elements in the window.
func AvgPool3D(input *core.Array, kernel [3]int, stride [3]int, padding [3]int, countIncludePad bool) (*core.Array, error) {
	return pool3D(input, kernel, stride, padding, false, countIncludePad)
}

// AdaptiveMaxPool2D pools to the requested output size by taking the
// maximum over each output window. Window o covers the input rows
// from floor(o·H_in/H_out) to ceil((o+1)·H_in/H_out), so every input
// sample lands in some window. A NaN in a window propagates: the
// maximum answers NaN.
func AdaptiveMaxPool2D(input *core.Array, outputH, outputW int) (*core.Array, error) {
	if input.NDim() != 4 {
		return nil, base.Errf("AdaptiveMaxPool2D: input must be 4-D, got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("AdaptiveMaxPool2D: complex arrays are not supported")
	}
	if outputH < 1 || outputW < 1 {
		return nil, base.Errf("AdaptiveMaxPool2D: output size must be at least 1")
	}
	n, c, hIn, wIn := input.Shape()[0], input.Shape()[1], input.Shape()[2], input.Shape()[3]
	// A zero spatial dimension leaves every window empty: max would
	// answer -Inf, so it is an error like every other shape refusal.
	if hIn < 1 || wIn < 1 {
		return nil, base.Errf("AdaptiveMaxPool2D: the spatial dimensions must not be empty, got shape %s", base.ShapeText(input.Shape()))
	}
	out := core.New(core.Float, []int{n, c, outputH, outputW}...)
	inF := input.RawFloats()
	if inF == nil || input.Strided() {
		inF = widenFloats(input)
	}
	outF := out.RawFloats()
	// One work item per output row; rows are disjoint. A row walks
	// about one window height of input rows per window column.
	items := n * c * outputH
	floor := workFloorFor(max(1, (hIn+outputH-1)/outputH) * wIn)
	engine.ParallelMin(items, floor, func(bs, be int) {
		for item := bs; item < be; item++ {
			oy := item % outputH
			row := item / outputH
			ch := row % c
			batch := row / c
			chanBase := batch*c*hIn*wIn + ch*hIn*wIn
			yStart := oy * hIn / outputH
			yEnd := ((oy+1)*hIn + outputH - 1) / outputH
			for ox := range outputW {
				xStart := ox * wIn / outputW
				xEnd := ((ox+1)*wIn + outputW - 1) / outputW
				best := math.Inf(-1)
				for y := yStart; y < yEnd; y++ {
					inRow := chanBase + y*wIn
					for x := xStart; x < xEnd; x++ {
						if v := inF[inRow+x]; v > best || math.IsNaN(v) {
							best = v
						}
					}
				}
				outF[batch*c*outputH*outputW+ch*outputH*outputW+oy*outputW+ox] = best
			}
		}
	})
	return out, nil
}

// GlobalAvgPool2D returns AdaptiveAvgPool2D reduced to (1, 1), the
// global average pooling operation common in modern CNNs.
func GlobalAvgPool2D(input *core.Array) (*core.Array, error) {
	return AdaptiveAvgPool2D(input, 1, 1)
}

// GlobalMaxPool2D returns the global maximum of a 4-D feature map.
func GlobalMaxPool2D(input *core.Array) (*core.Array, error) {
	return AdaptiveMaxPool2D(input, 1, 1)
}

// GlobalAvgPool1D reduces a 3-D input (N, C, L) to shape (N, C, 1).
func GlobalAvgPool1D(input *core.Array) (*core.Array, error) {
	if input.NDim() != 3 {
		return nil, base.Errf("GlobalAvgPool1D: input must be 3-D (N, C, L), got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("GlobalAvgPool1D: complex arrays are not supported")
	}
	if input.Shape()[2] == 0 {
		return nil, base.Errf("GlobalAvgPool1D: the spatial dimension must not be empty")
	}
	return pool1D(input, input.Shape()[2], 1, 0, false, false)
}

// GlobalMaxPool1D reduces a 3-D input (N, C, L) to shape (N, C, 1).
func GlobalMaxPool1D(input *core.Array) (*core.Array, error) {
	if input.NDim() != 3 {
		return nil, base.Errf("GlobalMaxPool1D: input must be 3-D (N, C, L), got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("GlobalMaxPool1D: complex arrays are not supported")
	}
	if input.Shape()[2] == 0 {
		return nil, base.Errf("GlobalMaxPool1D: the spatial dimension must not be empty")
	}
	return pool1D(input, input.Shape()[2], 1, 0, true, false)
}

// GlobalAvgPool3D reduces a 5-D input (N, C, D, H, W) to (N, C, 1, 1, 1).
func GlobalAvgPool3D(input *core.Array) (*core.Array, error) {
	if input.NDim() != 5 {
		return nil, base.Errf("GlobalAvgPool3D: input must be 5-D (N, C, D, H, W), got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("GlobalAvgPool3D: complex arrays are not supported")
	}
	if input.Shape()[2] == 0 || input.Shape()[3] == 0 || input.Shape()[4] == 0 {
		return nil, base.Errf("GlobalAvgPool3D: the spatial dimensions must not be empty")
	}
	return pool3D(input, [3]int{input.Shape()[2], input.Shape()[3], input.Shape()[4]}, [3]int{1, 1, 1}, [3]int{0, 0, 0}, false, false)
}

// GlobalMaxPool3D reduces a 5-D input (N, C, D, H, W) to (N, C, 1, 1, 1).
func GlobalMaxPool3D(input *core.Array) (*core.Array, error) {
	if input.NDim() != 5 {
		return nil, base.Errf("GlobalMaxPool3D: input must be 5-D (N, C, D, H, W), got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("GlobalMaxPool3D: complex arrays are not supported")
	}
	if input.Shape()[2] == 0 || input.Shape()[3] == 0 || input.Shape()[4] == 0 {
		return nil, base.Errf("GlobalMaxPool3D: the spatial dimensions must not be empty")
	}
	return pool3D(input, [3]int{input.Shape()[2], input.Shape()[3], input.Shape()[4]}, [3]int{1, 1, 1}, [3]int{0, 0, 0}, true, false)
}

// AdaptiveMaxPool1D pools the input to the requested output size by
// taking the maximum over each output window. Window o covers the
// input positions from floor(o·L_in/L_out) to ceil((o+1)·L_in/L_out),
// so every input sample lands in some window. Input is (N, C, L). A
// NaN in a window propagates: the maximum answers NaN.
func AdaptiveMaxPool1D(input *core.Array, outputL int) (*core.Array, error) {
	if input.NDim() != 3 {
		return nil, base.Errf("AdaptiveMaxPool1D: input must be 3-D, got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("AdaptiveMaxPool1D: complex arrays are not supported")
	}
	if outputL < 1 {
		return nil, base.Errf("AdaptiveMaxPool1D: output size must be at least 1")
	}
	n, c, lIn := input.Shape()[0], input.Shape()[1], input.Shape()[2]
	if lIn < 1 {
		return nil, base.Errf("AdaptiveMaxPool1D: the spatial dimension must not be empty, got shape %s", base.ShapeText(input.Shape()))
	}
	out := core.New(core.Float, []int{n, c, outputL}...)
	inF := input.RawFloats()
	if inF == nil || input.Strided() {
		inF = widenFloats(input)
	}
	outF := out.RawFloats()
	// One work item per (batch, channel) row; rows are disjoint.
	engine.ParallelMin(n*c, workFloorFor(lIn), func(bs, be int) {
		for row := bs; row < be; row++ {
			ch := row % c
			batch := row / c
			chanBase := batch*c*lIn + ch*lIn
			for ol := range outputL {
				start := ol * lIn / outputL
				end := ((ol+1)*lIn + outputL - 1) / outputL
				best := math.Inf(-1)
				for l := start; l < end; l++ {
					if v := inF[chanBase+l]; v > best || math.IsNaN(v) {
						best = v
					}
				}
				outF[batch*c*outputL+ch*outputL+ol] = best
			}
		}
	})
	return out, nil
}

// AdaptiveMaxPool3D pools the input to the requested output size by
// taking the maximum over each output window, with the same
// floor-start, ceil-end convention as AdaptiveMaxPool1D. Input is
// (N, C, D, H, W). A NaN in a window propagates: the maximum answers
// NaN.
func AdaptiveMaxPool3D(input *core.Array, outputD, outputH, outputW int) (*core.Array, error) {
	if input.NDim() != 5 {
		return nil, base.Errf("AdaptiveMaxPool3D: input must be 5-D, got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("AdaptiveMaxPool3D: complex arrays are not supported")
	}
	if outputD < 1 || outputH < 1 || outputW < 1 {
		return nil, base.Errf("AdaptiveMaxPool3D: output sizes must be at least 1")
	}
	n, c, dIn, hIn, wIn := input.Shape()[0], input.Shape()[1], input.Shape()[2], input.Shape()[3], input.Shape()[4]
	if dIn < 1 || hIn < 1 || wIn < 1 {
		return nil, base.Errf("AdaptiveMaxPool3D: the spatial dimensions must not be empty, got shape %s", base.ShapeText(input.Shape()))
	}
	out := core.New(core.Float, []int{n, c, outputD, outputH, outputW}...)
	strideHW := wIn
	strideDHW := hIn * wIn
	strideCDHW := dIn * hIn * wIn
	inF := input.RawFloats()
	if inF == nil || input.Strided() {
		inF = widenFloats(input)
	}
	outF := out.RawFloats()
	// One work item per output depth-height row; rows are disjoint.
	items := n * c * outputD * outputH
	floor := workFloorFor(wIn * max(1, (hIn+outputH-1)/outputH) * max(1, (dIn+outputD-1)/outputD))
	engine.ParallelMin(items, floor, func(bs, be int) {
		for item := bs; item < be; item++ {
			oh := item % outputH
			t := item / outputH
			od := t % outputD
			t /= outputD
			ch := t % c
			batch := t / c
			chanBase := batch*c*strideCDHW + ch*strideCDHW
			dStart := od * dIn / outputD
			dEnd := ((od+1)*dIn + outputD - 1) / outputD
			hStart := oh * hIn / outputH
			hEnd := ((oh+1)*hIn + outputH - 1) / outputH
			for ow := range outputW {
				wStart := ow * wIn / outputW
				wEnd := ((ow+1)*wIn + outputW - 1) / outputW
				best := math.Inf(-1)
				for d := dStart; d < dEnd; d++ {
					for h := hStart; h < hEnd; h++ {
						inRow := chanBase + d*strideDHW + h*strideHW
						for w := wStart; w < wEnd; w++ {
							if v := inF[inRow+w]; v > best || math.IsNaN(v) {
								best = v
							}
						}
					}
				}
				outF[batch*c*outputD*outputH*outputW+ch*outputD*outputH*outputW+od*outputH*outputW+oh*outputW+ow] = best
			}
		}
	})
	return out, nil
}

// AdaptiveAvgPool1D pools to the requested output size by averaging
// each output window, with the same floor-start, ceil-end convention
// as AdaptiveMaxPool1D. Input is (N, C, L).
func AdaptiveAvgPool1D(input *core.Array, outputL int) (*core.Array, error) {
	if input.NDim() != 3 {
		return nil, base.Errf("AdaptiveAvgPool1D: input must be 3-D, got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("AdaptiveAvgPool1D: complex arrays are not supported")
	}
	if outputL < 1 {
		return nil, base.Errf("AdaptiveAvgPool1D: output size must be at least 1")
	}
	n, c, lIn := input.Shape()[0], input.Shape()[1], input.Shape()[2]
	if lIn < 1 {
		return nil, base.Errf("AdaptiveAvgPool1D: the spatial dimension must not be empty, got shape %s", base.ShapeText(input.Shape()))
	}
	out := core.New(core.Float, []int{n, c, outputL}...)
	inF := input.RawFloats()
	if inF == nil || input.Strided() {
		inF = widenFloats(input)
	}
	outF := out.RawFloats()
	// One work item per (batch, channel) row; rows are disjoint and
	// every window sums its elements in ascending order.
	engine.ParallelMin(n*c, workFloorFor(lIn), func(bs, be int) {
		for row := bs; row < be; row++ {
			ch := row % c
			batch := row / c
			chanBase := batch*c*lIn + ch*lIn
			for ol := range outputL {
				start := ol * lIn / outputL
				end := ((ol+1)*lIn + outputL - 1) / outputL
				var sum float64
				for l := start; l < end; l++ {
					sum += inF[chanBase+l]
				}
				outF[batch*c*outputL+ch*outputL+ol] = sum / float64(end-start)
			}
		}
	})
	return out, nil
}

// AdaptiveAvgPool3D pools to the requested output size by averaging
// each output window, with the same floor-start, ceil-end convention
// as AdaptiveMaxPool1D. Input is (N, C, D, H, W).
func AdaptiveAvgPool3D(input *core.Array, outputD, outputH, outputW int) (*core.Array, error) {
	if input.NDim() != 5 {
		return nil, base.Errf("AdaptiveAvgPool3D: input must be 5-D, got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("AdaptiveAvgPool3D: complex arrays are not supported")
	}
	if outputD < 1 || outputH < 1 || outputW < 1 {
		return nil, base.Errf("AdaptiveAvgPool3D: output sizes must be at least 1")
	}
	n, c, dIn, hIn, wIn := input.Shape()[0], input.Shape()[1], input.Shape()[2], input.Shape()[3], input.Shape()[4]
	if dIn < 1 || hIn < 1 || wIn < 1 {
		return nil, base.Errf("AdaptiveAvgPool3D: the spatial dimensions must not be empty, got shape %s", base.ShapeText(input.Shape()))
	}
	out := core.New(core.Float, []int{n, c, outputD, outputH, outputW}...)
	strideHW := wIn
	strideDHW := hIn * wIn
	strideCDHW := dIn * hIn * wIn
	inF := input.RawFloats()
	if inF == nil || input.Strided() {
		inF = widenFloats(input)
	}
	outF := out.RawFloats()
	// One work item per output element here: the 3-D windows share no
	// row structure worth exploiting, and every element sums its own
	// window in ascending order.
	items := n * c * outputD * outputH * outputW
	floor := workFloorFor(max(1, (dIn+outputD-1)/outputD) * max(1, (hIn+outputH-1)/outputH) * max(1, (wIn+outputW-1)/outputW))
	engine.ParallelMin(items, floor, func(bs, be int) {
		for item := bs; item < be; item++ {
			ow := item % outputW
			t := item / outputW
			oh := t % outputH
			t /= outputH
			od := t % outputD
			t /= outputD
			ch := t % c
			batch := t / c
			chanBase := batch*c*strideCDHW + ch*strideCDHW
			dStart := od * dIn / outputD
			dEnd := ((od+1)*dIn + outputD - 1) / outputD
			hStart := oh * hIn / outputH
			hEnd := ((oh+1)*hIn + outputH - 1) / outputH
			wStart := ow * wIn / outputW
			wEnd := ((ow+1)*wIn + outputW - 1) / outputW
			var sum float64
			count := 0
			for d := dStart; d < dEnd; d++ {
				for h := hStart; h < hEnd; h++ {
					inRow := chanBase + d*strideDHW + h*strideHW
					for w := wStart; w < wEnd; w++ {
						sum += inF[inRow+w]
						count++
					}
				}
			}
			if count > 0 {
				outF[batch*c*outputD*outputH*outputW+ch*outputD*outputH*outputW+od*outputH*outputW+oh*outputW+ow] = sum / float64(count)
			}
		}
	})
	return out, nil
}

// AdaptiveAvgPool2D pools the input to the requested output size by
// averaging each output window, with the same floor-start, ceil-end
// convention as AdaptiveMaxPool2D. outputH and outputW must both be
// at least 1.
func AdaptiveAvgPool2D(input *core.Array, outputH, outputW int) (*core.Array, error) {
	if input.NDim() != 4 {
		return nil, base.Errf("AdaptiveAvgPool2D: input must be 4-D, got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("AdaptiveAvgPool2D: complex arrays are not supported")
	}
	if outputH < 1 || outputW < 1 {
		return nil, base.Errf("AdaptiveAvgPool2D: output size must be at least 1, got %dx%d", outputH, outputW)
	}
	n, c, hIn, wIn := input.Shape()[0], input.Shape()[1], input.Shape()[2], input.Shape()[3]
	// A zero spatial dimension leaves every window empty: avg would
	// divide by zero, so it is an error like every other shape refusal.
	if hIn < 1 || wIn < 1 {
		return nil, base.Errf("AdaptiveAvgPool2D: the spatial dimensions must not be empty, got shape %s", base.ShapeText(input.Shape()))
	}
	out := core.New(core.Float, []int{n, c, outputH, outputW}...)
	inF := input.RawFloats()
	if inF == nil || input.Strided() {
		inF = widenFloats(input)
	}
	outF := out.RawFloats()
	// One work item per output row; rows are disjoint and every window
	// sums its elements in ascending order.
	items := n * c * outputH
	engine.ParallelMin(items, workFloorFor(max(1, (hIn+outputH-1)/outputH)*wIn), func(bs, be int) {
		for item := bs; item < be; item++ {
			oy := item % outputH
			row := item / outputH
			ch := row % c
			batch := row / c
			chanBase := batch*c*hIn*wIn + ch*hIn*wIn
			// Window in the input: floor start, ceil end, so
			// every input sample lands in some window.
			yStart := oy * hIn / outputH
			yEnd := ((oy+1)*hIn + outputH - 1) / outputH
			for ox := range outputW {
				xStart := ox * wIn / outputW
				xEnd := ((ox+1)*wIn + outputW - 1) / outputW
				var sum float64
				for y := yStart; y < yEnd; y++ {
					inRow := chanBase + y*wIn
					for x := xStart; x < xEnd; x++ {
						sum += inF[inRow+x]
					}
				}
				count := float64((yEnd - yStart) * (xEnd - xStart))
				off := batch*c*outputH*outputW + ch*outputH*outputW + oy*outputW + ox
				outF[off] = sum / count
			}
		}
	})
	return out, nil
}

// pool2D is the shared implementation behind MaxPool2D and AvgPool2D.
func pool2D(input *core.Array, kernel, stride, padding int, isMax, countIncludePad bool) (*core.Array, error) {
	if input.NDim() != 4 {
		return nil, base.Errf("Pool2D: input must be 4-D, got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("Pool2D: complex arrays are not supported")
	}
	if kernel < 1 {
		return nil, base.Errf("Pool2D: kernel must be at least 1, got %d", kernel)
	}
	if stride < 1 {
		return nil, base.Errf("Pool2D: stride must be at least 1, got %d", stride)
	}
	if padding < 0 {
		return nil, base.Errf("Pool2D: padding must be non-negative, got %d", padding)
	}
	// Padding from the kernel upward leaves output windows entirely
	// inside the padding (the first window covers −padding ..
	// −padding+kernel−1): a max would answer −Inf and an average 0,
	// so such configurations are refused, in the spirit of the
	// common pooling implementations' padding bound.
	if padding >= kernel {
		return nil, base.Errf("Pool2D: padding %d must stay below the kernel %d, some windows would hold no data", padding, kernel)
	}
	n, c, hIn, wIn := input.Shape()[0], input.Shape()[1], input.Shape()[2], input.Shape()[3]
	// Guard the numerators before the division: a negative numerator
	// truncates toward zero and masquerades as a 1-wide output.
	hNum := hIn + 2*padding - kernel
	wNum := wIn + 2*padding - kernel
	if hNum < 0 || wNum < 0 {
		return nil, base.Errf("Pool2D: kernel %d with padding %d does not fit the %dx%d input", kernel, padding, hIn, wIn)
	}
	hOut := hNum/stride + 1
	wOut := wNum/stride + 1
	if hOut < 1 || wOut < 1 {
		return nil, base.Errf("Pool2D: output is empty (H_out=%d, W_out=%d)", hOut, wOut)
	}
	out := core.New(core.Float, []int{n, c, hOut, wOut}...)
	inF := input.RawFloats()
	if inF == nil || input.Strided() {
		inF = widenFloats(input)
	}
	outF := out.RawFloats()
	// One work item per output row (batch, channel, oy). Rows are
	// disjoint. The max and average variants walk the same windows in
	// the same (kernel height, kernel width) order, but as two loop
	// bodies: the max/sum choice leaves the innermost loop, which would
	// otherwise pay for a perfectly predicted branch on every tap, and
	// neither variant's arithmetic changes a bit.
	items := n * c * hOut
	engine.ParallelMin(items, workFloorFor(wOut*kernel*kernel), func(bs, be int) {
		if isMax {
			for item := bs; item < be; item++ {
				oy := item % hOut
				row := item / hOut
				ch := row % c
				batch := row / c
				chanBase := batch*c*hIn*wIn + ch*hIn*wIn
				for ox := range wOut {
					best := math.Inf(-1)
					for kh := range kernel {
						iy := oy*stride + kh - padding
						if iy < 0 || iy >= hIn {
							continue
						}
						inRow := chanBase + iy*wIn
						for kw := range kernel {
							ix := ox*stride + kw - padding
							if ix < 0 || ix >= wIn {
								continue
							}
							if v := inF[inRow+ix]; v > best || math.IsNaN(v) {
								best = v
							}
						}
					}
					off := batch*c*hOut*wOut + ch*hOut*wOut + oy*wOut + ox
					outF[off] = best
				}
			}
			return
		}
		for item := bs; item < be; item++ {
			oy := item % hOut
			row := item / hOut
			ch := row % c
			batch := row / c
			chanBase := batch*c*hIn*wIn + ch*hIn*wIn
			for ox := range wOut {
				var acc float64
				var count int
				for kh := range kernel {
					iy := oy*stride + kh - padding
					if iy < 0 || iy >= hIn {
						continue
					}
					inRow := chanBase + iy*wIn
					for kw := range kernel {
						ix := ox*stride + kw - padding
						if ix < 0 || ix >= wIn {
							continue
						}
						acc += inF[inRow+ix]
						count++
					}
				}
				if countIncludePad {
					acc /= float64(kernel * kernel)
				} else if count > 0 {
					acc /= float64(count)
				}
				off := batch*c*hOut*wOut + ch*hOut*wOut + oy*wOut + ox
				outF[off] = acc
			}
		}
	})
	return out, nil
}

// pool1D is the shared implementation behind MaxPool1D and AvgPool1D.
// countIncludePad is honoured only by the average variant.
func pool1D(input *core.Array, kernel, stride, padding int, isMax, countIncludePad bool) (*core.Array, error) {
	if input.NDim() != 3 {
		return nil, base.Errf("Pool1D: input must be 3-D, got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("Pool1D: complex arrays are not supported")
	}
	if kernel < 1 || stride < 1 || padding < 0 {
		return nil, base.Errf("Pool1D: kernel/stride ≥1, padding ≥0")
	}
	// Padding from the kernel upward leaves output windows entirely
	// inside the padding (the first window covers −padding ..
	// −padding+kernel−1): a max would answer −Inf and an average 0,
	// so such configurations are refused, in the spirit of the
	// common pooling implementations' padding bound.
	if padding >= kernel {
		return nil, base.Errf("Pool1D: padding %d must stay below the kernel %d, some windows would hold no data", padding, kernel)
	}
	n, c, lIn := input.Shape()[0], input.Shape()[1], input.Shape()[2]
	lNum := lIn + 2*padding - kernel
	if lNum < 0 {
		return nil, base.Errf("Pool1D: kernel %d with padding %d does not fit the length-%d input", kernel, padding, lIn)
	}
	lOut := lNum/stride + 1
	if lOut < 1 {
		return nil, base.Errf("Pool1D: output is empty (L_out=%d)", lOut)
	}
	out := core.New(core.Float, []int{n, c, lOut}...)
	inF := input.RawFloats()
	if inF == nil || input.Strided() {
		inF = widenFloats(input)
	}
	outF := out.RawFloats()
	// One work item per (batch, channel) row; rows are disjoint and
	// every window walks its elements in ascending order.
	engine.ParallelMin(n*c, workFloorFor(lOut*kernel), func(bs, be int) {
		for row := bs; row < be; row++ {
			ch := row % c
			batch := row / c
			chanBase := batch*c*lIn + ch*lIn
			for ol := range lOut {
				var acc float64
				var best float64
				if isMax {
					best = math.Inf(-1)
				}
				var count int
				for kl := range kernel {
					il := ol*stride + kl - padding
					if il < 0 || il >= lIn {
						continue
					}
					v := inF[chanBase+il]
					if isMax {
						if v > best || math.IsNaN(v) {
							best = v
						}
					} else {
						acc += v
					}
					count++
				}
				if isMax {
					outF[batch*c*lOut+ch*lOut+ol] = best
				} else if countIncludePad {
					outF[batch*c*lOut+ch*lOut+ol] = acc / float64(kernel)
				} else if count > 0 {
					outF[batch*c*lOut+ch*lOut+ol] = acc / float64(count)
				}
			}
		}
	})
	return out, nil
}

// pool3D is the shared implementation behind MaxPool3D and AvgPool3D.
// countIncludePad is honoured only by the average variant.
func pool3D(input *core.Array, kernel, stride, padding [3]int, isMax, countIncludePad bool) (*core.Array, error) {
	if input.NDim() != 5 {
		return nil, base.Errf("Pool3D: input must be 5-D, got shape %s", base.ShapeText(input.Shape()))
	}
	if input.Dtype() == core.Complex {
		return nil, base.Errf("Pool3D: complex arrays are not supported")
	}
	for d := range kernel {
		if kernel[d] < 1 || stride[d] < 1 || padding[d] < 0 {
			return nil, base.Errf("Pool3D: kernel/stride ≥1, padding ≥0")
		}
		// Padding from the kernel upward leaves output windows
		// entirely inside the padding: a max would answer −Inf and
		// an average 0, so such configurations are refused, in the
		// spirit of the common pooling implementations' padding
		// bound.
		if padding[d] >= kernel[d] {
			return nil, base.Errf("Pool3D: padding %d must stay below the kernel %d in dimension %d, some windows would hold no data", padding[d], kernel[d], d)
		}
	}
	n, c, dIn, hIn, wIn := input.Shape()[0], input.Shape()[1], input.Shape()[2], input.Shape()[3], input.Shape()[4]
	// Guard the numerators before the division: a negative numerator
	// truncates toward zero and masquerades as a 1-wide output.
	dNum := dIn + 2*padding[0] - kernel[0]
	hNum := hIn + 2*padding[1] - kernel[1]
	wNum := wIn + 2*padding[2] - kernel[2]
	if dNum < 0 || hNum < 0 || wNum < 0 {
		return nil, base.Errf("Pool3D: kernel %v with padding %v does not fit the input", kernel, padding)
	}
	dOut := dNum/stride[0] + 1
	hOut := hNum/stride[1] + 1
	wOut := wNum/stride[2] + 1
	if dOut < 1 || hOut < 1 || wOut < 1 {
		return nil, base.Errf("Pool3D: output is empty")
	}
	out := core.New(core.Float, []int{n, c, dOut, hOut, wOut}...)
	strideHW := wIn
	strideDHW := hIn * wIn
	strideCDHW := dIn * hIn * wIn
	inF := input.RawFloats()
	if inF == nil || input.Strided() {
		inF = widenFloats(input)
	}
	outF := out.RawFloats()
	// One work item per output depth-height row (batch, channel, od,
	// oh); rows are disjoint and each window walks its elements in the
	// (kernel depth, kernel height, kernel width) order.
	items := n * c * dOut * hOut
	engine.ParallelMin(items, workFloorFor(wOut*kernel[0]*kernel[1]*kernel[2]), func(bs, be int) {
		for item := bs; item < be; item++ {
			oh := item % hOut
			t := item / hOut
			od := t % dOut
			t /= dOut
			ch := t % c
			batch := t / c
			chanBase := batch*c*strideCDHW + ch*strideCDHW
			for ow := range wOut {
				var acc float64
				var best float64
				if isMax {
					best = math.Inf(-1)
				}
				var count int
				for kd := range kernel[0] {
					id := od*stride[0] + kd - padding[0]
					if id < 0 || id >= dIn {
						continue
					}
					inDepth := chanBase + id*strideDHW
					for kh := range kernel[1] {
						ih := oh*stride[1] + kh - padding[1]
						if ih < 0 || ih >= hIn {
							continue
						}
						inRow := inDepth + ih*strideHW
						for kw := range kernel[2] {
							iw := ow*stride[2] + kw - padding[2]
							if iw < 0 || iw >= wIn {
								continue
							}
							v := inF[inRow+iw]
							if isMax {
								if v > best || math.IsNaN(v) {
									best = v
								}
							} else {
								acc += v
							}
							count++
						}
					}
				}
				off := batch*c*dOut*hOut*wOut + ch*dOut*hOut*wOut + od*hOut*wOut + oh*wOut + ow
				if isMax {
					outF[off] = best
				} else if countIncludePad {
					outF[off] = acc / float64(kernel[0]*kernel[1]*kernel[2])
				} else if count > 0 {
					outF[off] = acc / float64(count)
				}
			}
		}
	})
	return out, nil
}