// Copyright (C) 2026 Kiyotsugu Arai
// SPDX-License-Identifier: LGPL-3.0-or-later

// fft2d.hpp — 2-dimensional FFT / IFFT
//
// New implementation:
//   - Realizes 2D FFT by applying 1D FFT along rows then along columns
//   - Matrix<complex> based
//   - Uses sangi::FFT internally

#ifndef SANGI_FFT_2D_HPP
#define SANGI_FFT_2D_HPP

#include <math/fft/fft.hpp>
#include <math/core/matrix.hpp>
#include <math/core/Complex.hpp>
#include <vector>

namespace sangi {

    /// 2D FFT (along rows → along columns)
    /// @param input  input matrix (rows × cols, complex)
    /// @return       frequency-domain matrix (same size)
    template<typename Real = double>
    Matrix<Complex<Real>> fft2d(const Matrix<Complex<Real>>& input) {
        using C = Complex<Real>;
        std::size_t rows = input.rows();
        std::size_t cols = input.cols();

        Matrix<C> result = input;

        // FFT along rows
        std::vector<C> rowBuf(cols);
        for (std::size_t r = 0; r < rows; ++r) {
            for (std::size_t c = 0; c < cols; ++c) rowBuf[c] = result(r, c);
            FFT<C, Real>::fft(std::span<C>(rowBuf));
            for (std::size_t c = 0; c < cols; ++c) result(r, c) = rowBuf[c];
        }

        // FFT along columns
        std::vector<C> colBuf(rows);
        for (std::size_t c = 0; c < cols; ++c) {
            for (std::size_t r = 0; r < rows; ++r) colBuf[r] = result(r, c);
            FFT<C, Real>::fft(std::span<C>(colBuf));
            for (std::size_t r = 0; r < rows; ++r) result(r, c) = colBuf[r];
        }

        return result;
    }

    /// 2D IFFT (along columns → along rows)
    /// @param input  frequency-domain matrix
    /// @return       spatial-domain matrix
    template<typename Real = double>
    Matrix<Complex<Real>> ifft2d(const Matrix<Complex<Real>>& input) {
        using C = Complex<Real>;
        std::size_t rows = input.rows();
        std::size_t cols = input.cols();

        Matrix<C> result = input;

        // IFFT along rows
        std::vector<C> rowBuf(cols);
        for (std::size_t r = 0; r < rows; ++r) {
            for (std::size_t c = 0; c < cols; ++c) rowBuf[c] = result(r, c);
            FFT<C, Real>::ifft(std::span<C>(rowBuf));
            for (std::size_t c = 0; c < cols; ++c) result(r, c) = rowBuf[c];
        }

        // IFFT along columns
        std::vector<C> colBuf(rows);
        for (std::size_t c = 0; c < cols; ++c) {
            for (std::size_t r = 0; r < rows; ++r) colBuf[r] = result(r, c);
            FFT<C, Real>::ifft(std::span<C>(colBuf));
            for (std::size_t r = 0; r < rows; ++r) result(r, c) = colBuf[r];
        }

        return result;
    }

    /// Real matrix → 2D FFT
    /// @param input  input matrix (real)
    /// @return       frequency-domain complex matrix
    template<typename Real = double>
    Matrix<Complex<Real>> fft2d_real(const Matrix<Real>& input) {
        using C = Complex<Real>;
        std::size_t rows = input.rows();
        std::size_t cols = input.cols();

        Matrix<C> complexInput(rows, cols);
        for (std::size_t r = 0; r < rows; ++r)
            for (std::size_t c = 0; c < cols; ++c)
                complexInput(r, c) = C(input(r, c), Real{0});

        return fft2d(complexInput);
    }

} // namespace sangi

#endif // SANGI_FFT_2D_HPP