from __future__ import annotations
import numpy as np


def truncated_fourier_weights(mask: np.ndarray, bandwidth, *, clip_negative=False) -> np.ndarray:
    """Fourier-project a 2-D or 3-D periodic grid indicator."""
    mask = np.asarray(mask, dtype=float)
    if mask.ndim not in (2, 3):
        raise ValueError("mask must be 2-D or 3-D")
    if len(bandwidth) != mask.ndim:
        raise ValueError("bandwidth must have one cutoff per axis")
    spectrum = np.fft.fftn(mask)
    keep = np.ones(mask.shape, dtype=bool)
    for axis, cutoff in enumerate(bandwidth):
        freq = np.fft.fftfreq(mask.shape[axis]) * mask.shape[axis]
        shape = [1] * mask.ndim
        shape[axis] = mask.shape[axis]
        keep &= np.abs(freq).reshape(shape) <= int(cutoff)
    weights = np.fft.ifftn(spectrum * keep).real
    return np.maximum(weights, 0.0) if clip_negative else weights


def pool(features: np.ndarray, weights: np.ndarray, eps=1e-12) -> np.ndarray:
    """Normalized weighted pooling for C,*grid or B,C,*grid."""
    x, w = np.asarray(features), np.asarray(weights)
    if x.shape[-w.ndim:] != w.shape or x.ndim not in (w.ndim + 1, w.ndim + 2):
        raise ValueError("features must be C,*grid or B,C,*grid")
    axes = tuple(range(x.ndim - w.ndim, x.ndim))
    return np.sum(x * w, axis=axes) / (np.sum(w) + eps)