# ndwasm_signal.js

    import { NDArray } from './ndarray_core.js';
    import { NDWasm, fromWasm } from './ndwasm.js';
    /**
     * NDWasmSignal: Signal Processing & Transformations
     * Handles O(n log n) frequency domain transforms and O(n^2 * k^2) spatial filters.
     * @namespace NDWasmSignal
     */
    export const NDWasmSignal = {
        /**
         * 1D Complex-to-Complex Fast Fourier Transform.
         * The input array must have its last dimension of size 2 (real and imaginary parts).
         * The transform is performed in-place.
         * Complexity: O(n log n)
         * @memberof NDWasmSignal
         * @param {NDArray} a - Complex input signal, with shape [..., 2].
         * @returns {NDArray} - Complex result, with the same shape as input.
         */
        fft(a) {
            if (a.ndim !== 2 || a.shape[1] !== 2) {
                throw new Error("Input to fft must be a 1D complex array with shape [n, 2].");
            }

            const n = a.size / 2;
            const suffix = NDWasm.runtime._getSuffix(a.dtype);
            const wComplex = a.toWasm(NDWasm.runtime);

            try {
                NDWasm.runtime.exports[`FFT1D${suffix}`](wComplex.ptr, n);
                return fromWasm(wComplex, a.shape, a.dtype);
            } finally {
                wComplex.dispose();
            }
        },

        /**
         * 1D Inverse Complex-to-Complex Fast Fourier Transform.
         * The input array must have its last dimension of size 2 (real and imaginary parts).
         * The transform is performed in-place.
         * Complexity: O(n log n)
         * @memberof NDWasmSignal
         * @param {NDArray} a - Complex frequency-domain signal, with shape [..., 2].
         * @returns {NDArray} - Complex time-domain result, with the same shape as input.
         */
        ifft(a) {
            if (a.ndim !== 2 || a.shape[1] !== 2) {
                throw new Error("Input to ifft must be a 1D complex array with shape [n, 2].");
            }

            const n = a.size / 2;
            const suffix = NDWasm.runtime._getSuffix(a.dtype);
            const wComplex = a.toWasm(NDWasm.runtime);

            try {
                NDWasm.runtime.exports[`IFFT1D${suffix}`](wComplex.ptr, n);
                return fromWasm(wComplex, a.shape, a.dtype);
            } finally {
                wComplex.dispose();
            }
        },

        /**
         * 1D Real-to-Complex Fast Fourier Transform (Optimized for real input).
         * The output is a complex array with shape [n/2 + 1, 2].
         * Complexity: O(n log n)
         * @memberof NDWasmSignal
         * @param {NDArray} a - Real input signal.
         * @returns {NDArray} - Complex result of shape [n/2 + 1, 2].
         */
        rfft(a) {
            if (a.ndim !== 1) {
                throw new Error("Input to rfft must be a 1D real array.");
            }
            const n = a.size;
            const outLen = Math.floor(n / 2) + 1;
            const suffix = NDWasm.runtime._getSuffix(a.dtype);
            const wa = a.toWasm(NDWasm.runtime);
            const wOut = NDWasm.runtime.createBuffer(outLen * 2, a.dtype);

            try {
                NDWasm.runtime.exports[`RFFT1D${suffix}`](wa.ptr, wOut.ptr, n);
                return fromWasm(wOut, [outLen, 2], a.dtype);
            } finally {
                wa.dispose();
                wOut.dispose();
            }
        },

        /**
         * 1D Complex-to-Real Inverse Fast Fourier Transform.
         * The input must be a complex array of shape [k, 2], where k is n/2 + 1.
         * @memberof NDWasmSignal
         * @param {NDArray} a - Complex frequency signal of shape [n/2 + 1, 2].
         * @param {number} n - Length of the original real signal.
         * @returns {NDArray} Real-valued time domain signal.
         */
        rifft(a, n) {
            if (a.ndim !== 2 || a.shape[1] !== 2) {
                throw new Error("Input to rifft must be a complex array with shape [k, 2].");
            }
            const suffix = NDWasm.runtime._getSuffix(a.dtype);
            const wa = a.toWasm(NDWasm.runtime);
            const wo = NDWasm.runtime.createBuffer(n, a.dtype);

            try {
                NDWasm.runtime.exports[`RIFFT1D${suffix}`](wa.ptr, wo.ptr, n);
                return fromWasm(wo, [n], a.dtype);
            } finally {
                wa.dispose();
                wo.dispose();
            }
        },

        /**
         * 2D Complex-to-Complex Fast Fourier Transform.
         * The input array must be 3D with shape [rows, cols, 2].
         * The transform is performed in-place.
         * Complexity: O(rows * cols * log(rows * cols))
         * @memberof NDWasmSignal
         * @param {NDArray} a - 2D Complex input signal, with shape [rows, cols, 2].
         * @returns {NDArray} - 2D Complex result, with the same shape as input.
         */
        fft2(a) {
            if (a.ndim !== 3 || a.shape[2] !== 2) {
                throw new Error("fft2 requires a 3D array with shape [rows, cols, 2].");
            }
            const [rows, cols] = a.shape;
            const suffix = NDWasm.runtime._getSuffix(a.dtype);
            const wComplex = a.toWasm(NDWasm.runtime);

            try {
                NDWasm.runtime.exports[`FFT2D${suffix}`](wComplex.ptr, rows, cols);
                return fromWasm(wComplex, a.shape, a.dtype);
            } finally {
                wComplex.dispose();
            }
        },

        /**
         * 2D Inverse Complex-to-Complex Fast Fourier Transform.
         * The input array must be 3D with shape [rows, cols, 2].
         * The transform is performed in-place.
         * @memberof NDWasmSignal
         * @param {NDArray} a - 2D Complex frequency-domain signal, with shape [rows, cols, 2].
         * @returns {NDArray} - 2D Complex time-domain result, with the same shape as input.
         */
        ifft2(a) {
            if (a.ndim !== 3 || a.shape[2] !== 2) {
                throw new Error("ifft2 requires a 3D array with shape [rows, cols, 2].");
            }
            const [rows, cols] = a.shape;
            const suffix = NDWasm.runtime._getSuffix(a.dtype);
            const wComplex = a.toWasm(NDWasm.runtime);

            try {
                NDWasm.runtime.exports[`IFFT2D${suffix}`](wComplex.ptr, rows, cols);
                return fromWasm(wComplex, a.shape, a.dtype);
            } finally {
                wComplex.dispose();
            }
        },

        /**
         * 1D Discrete Cosine Transform (Type II).
         * Complexity: O(n log n)
         * @memberof NDWasmSignal
         * @param {NDArray} a - Input signal.
         * @returns {NDArray} DCT result of same shape.
         */
        dct(a) {
            if (a.size < 2) {
                return a.copy();
            }
            const n = a.size;
            const suffix = NDWasm.runtime._getSuffix(a.dtype);
            return NDWasm._compute([a], a.shape, a.dtype, (aPtr, outPtr) => {
                return NDWasm.runtime.exports[`DCT${suffix}`](aPtr, outPtr, n);
            });
        },

        /**
         * 2D Spatial Convolution.
         * Complexity: O(img_h * img_w * kernel_h * kernel_w)
         * @memberof NDWasmSignal
         * @param {NDArray} img - 2D Image/Matrix.
         * @param {NDArray} kernel - 2D Filter kernel.
         * @param {number} stride - Step size (default 1).
         * @param {number} padding - Zero-padding size (default 0).
         * @returns {NDArray} Convolved result.
         */
        conv2d(img, kernel, stride = 1, padding = 0) {
            if (img.ndim !== 2 || kernel.ndim !== 2) throw new Error("Inputs must be 2D.");
            const [h, w] = img.shape;
            const [kh, kw] = kernel.shape;
            const oh = Math.floor((h - kh + 2 * padding) / stride) + 1;
            const ow = Math.floor((w - kw + 2 * padding) / stride) + 1;
            const suffix = NDWasm.runtime._getSuffix(img.dtype);

            return NDWasm._compute([img, kernel], [oh, ow], img.dtype, (iPtr, kPtr, outPtr) => {
                return NDWasm.runtime.exports[`Conv2D${suffix}`](iPtr, kPtr, outPtr, h, w, kh, kw, stride, padding);
            });
        },

        /**
         * 2D Spatial Cross-Correlation.
         * Similar to convolution but without flipping the kernel.
         * Complexity: O(img_h * img_w * kernel_h * kernel_w)
         * @memberof NDWasmSignal
         * @param {NDArray} img - 2D Image/Matrix.
         * @param {NDArray} kernel - 2D Filter kernel.
         * @param {number} stride - Step size.
         * @param {number} padding - Zero-padding size.
         * @returns {NDArray} Cross-correlated result.
         */
        correlate2d(img, kernel, stride = 1, padding = 0) {
            const [h, w] = img.shape;
            const [kh, kw] = kernel.shape;
            const oh = Math.floor((h - kh + 2 * padding) / stride) + 1;
            const ow = Math.floor((w - kw + 2 * padding) / stride) + 1;
            const suffix = NDWasm.runtime._getSuffix(img.dtype);

            return NDWasm._compute([img, kernel], [oh, ow], img.dtype, (iPtr, kPtr, outPtr) => {
                return NDWasm.runtime.exports[`CrossCorrelate2D${suffix}`](iPtr, kPtr, outPtr, h, w, kh, kw, stride, padding);
            });
        }
    };