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

// fft2d.hpp — 2 次元 FFT / IFFT
//
// 新規実装:
//   - 1D FFT の行方向→列方向適用で 2D FFT を実現
//   - Matrix<complex> ベース
//   - calx::FFT を内部で使用

#ifndef CALX_FFT_2D_HPP
#define CALX_FFT_2D_HPP

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

namespace calx {

    /// 2D FFT (行方向 → 列方向)
    /// @param input  入力行列 (rows × cols, 複素数)
    /// @return       周波数領域の行列 (同サイズ)
    template<typename Real = double>
    Matrix<std::complex<Real>> fft2d(const Matrix<std::complex<Real>>& input) {
        using C = std::complex<Real>;
        std::size_t rows = input.rows();
        std::size_t cols = input.cols();

        Matrix<C> result = input;

        // 行方向 FFT
        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
        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 (列方向 → 行方向)
    /// @param input  周波数領域の行列
    /// @return       空間領域の行列
    template<typename Real = double>
    Matrix<std::complex<Real>> ifft2d(const Matrix<std::complex<Real>>& input) {
        using C = std::complex<Real>;
        std::size_t rows = input.rows();
        std::size_t cols = input.cols();

        Matrix<C> result = input;

        // 行方向 IFFT
        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
        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;
    }

    /// 実数行列 → 2D FFT
    /// @param input  入力行列 (実数)
    /// @return       周波数領域の複素行列
    template<typename Real = double>
    Matrix<std::complex<Real>> fft2d_real(const Matrix<Real>& input) {
        using C = std::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 calx

#endif // CALX_FFT_2D_HPP