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

// fft3d.hpp — 3 次元 FFT / IFFT
//
// 新規実装:
//   - 1D FFT の x→y→z 方向逐次適用で 3D FFT を実現
//   - Tensor3D<complex> コンテナ + fft3d / ifft3d / fft3d_real
//   - calx::FFT を内部で使用

#ifndef CALX_FFT_3D_HPP
#define CALX_FFT_3D_HPP

#include <math/fft/fft.hpp>
#include <complex>
#include <vector>
#include <span>
#include <stdexcept>
#include <cmath>

namespace calx {

    // ── 3D テンソルコンテナ ──
    // row-major: data[z * ny*nx + y * nx + x]
    template<typename T>
    class Tensor3D {
    public:
        Tensor3D() : nx_(0), ny_(0), nz_(0) {}

        Tensor3D(std::size_t nx, std::size_t ny, std::size_t nz)
            : nx_(nx), ny_(ny), nz_(nz), data_(nx * ny * nz) {}

        Tensor3D(std::size_t nx, std::size_t ny, std::size_t nz, const T& val)
            : nx_(nx), ny_(ny), nz_(nz), data_(nx * ny * nz, val) {}

        std::size_t nx() const { return nx_; }
        std::size_t ny() const { return ny_; }
        std::size_t nz() const { return nz_; }
        std::size_t size() const { return data_.size(); }

        T& operator()(std::size_t x, std::size_t y, std::size_t z) {
            return data_[z * ny_ * nx_ + y * nx_ + x];
        }
        const T& operator()(std::size_t x, std::size_t y, std::size_t z) const {
            return data_[z * ny_ * nx_ + y * nx_ + x];
        }

        T* data() { return data_.data(); }
        const T* data() const { return data_.data(); }

    private:
        std::size_t nx_, ny_, nz_;
        std::vector<T> data_;
    };


    // ── 3D FFT (x→y→z 方向) ──
    template<typename Real = double>
    Tensor3D<std::complex<Real>> fft3d(const Tensor3D<std::complex<Real>>& input) {
        using C = std::complex<Real>;
        std::size_t nx = input.nx(), ny = input.ny(), nz = input.nz();

        Tensor3D<C> result = input;

        // x 方向 FFT
        {
            std::vector<C> buf(nx);
            for (std::size_t z = 0; z < nz; ++z) {
                for (std::size_t y = 0; y < ny; ++y) {
                    for (std::size_t x = 0; x < nx; ++x) buf[x] = result(x, y, z);
                    FFT<C, Real>::fft(std::span<C>(buf));
                    for (std::size_t x = 0; x < nx; ++x) result(x, y, z) = buf[x];
                }
            }
        }

        // y 方向 FFT
        {
            std::vector<C> buf(ny);
            for (std::size_t z = 0; z < nz; ++z) {
                for (std::size_t x = 0; x < nx; ++x) {
                    for (std::size_t y = 0; y < ny; ++y) buf[y] = result(x, y, z);
                    FFT<C, Real>::fft(std::span<C>(buf));
                    for (std::size_t y = 0; y < ny; ++y) result(x, y, z) = buf[y];
                }
            }
        }

        // z 方向 FFT
        {
            std::vector<C> buf(nz);
            for (std::size_t y = 0; y < ny; ++y) {
                for (std::size_t x = 0; x < nx; ++x) {
                    for (std::size_t z = 0; z < nz; ++z) buf[z] = result(x, y, z);
                    FFT<C, Real>::fft(std::span<C>(buf));
                    for (std::size_t z = 0; z < nz; ++z) result(x, y, z) = buf[z];
                }
            }
        }

        return result;
    }


    // ── 3D IFFT (x→y→z 方向) ──
    template<typename Real = double>
    Tensor3D<std::complex<Real>> ifft3d(const Tensor3D<std::complex<Real>>& input) {
        using C = std::complex<Real>;
        std::size_t nx = input.nx(), ny = input.ny(), nz = input.nz();

        Tensor3D<C> result = input;

        // x 方向 IFFT
        {
            std::vector<C> buf(nx);
            for (std::size_t z = 0; z < nz; ++z) {
                for (std::size_t y = 0; y < ny; ++y) {
                    for (std::size_t x = 0; x < nx; ++x) buf[x] = result(x, y, z);
                    FFT<C, Real>::ifft(std::span<C>(buf));
                    for (std::size_t x = 0; x < nx; ++x) result(x, y, z) = buf[x];
                }
            }
        }

        // y 方向 IFFT
        {
            std::vector<C> buf(ny);
            for (std::size_t z = 0; z < nz; ++z) {
                for (std::size_t x = 0; x < nx; ++x) {
                    for (std::size_t y = 0; y < ny; ++y) buf[y] = result(x, y, z);
                    FFT<C, Real>::ifft(std::span<C>(buf));
                    for (std::size_t y = 0; y < ny; ++y) result(x, y, z) = buf[y];
                }
            }
        }

        // z 方向 IFFT
        {
            std::vector<C> buf(nz);
            for (std::size_t y = 0; y < ny; ++y) {
                for (std::size_t x = 0; x < nx; ++x) {
                    for (std::size_t z = 0; z < nz; ++z) buf[z] = result(x, y, z);
                    FFT<C, Real>::ifft(std::span<C>(buf));
                    for (std::size_t z = 0; z < nz; ++z) result(x, y, z) = buf[z];
                }
            }
        }

        return result;
    }


    // ── 実数テンソル → 3D FFT ──
    template<typename Real = double>
    Tensor3D<std::complex<Real>> fft3d_real(const Tensor3D<Real>& input) {
        using C = std::complex<Real>;
        std::size_t nx = input.nx(), ny = input.ny(), nz = input.nz();

        Tensor3D<C> complexInput(nx, ny, nz);
        for (std::size_t z = 0; z < nz; ++z)
            for (std::size_t y = 0; y < ny; ++y)
                for (std::size_t x = 0; x < nx; ++x)
                    complexInput(x, y, z) = C(input(x, y, z), Real{0});

        return fft3d(complexInput);
    }


    // ── fftshift3d: DC 成分を中心に移動 ──
    template<typename T>
    Tensor3D<T> fftshift3d(const Tensor3D<T>& input) {
        std::size_t nx = input.nx(), ny = input.ny(), nz = input.nz();
        Tensor3D<T> out(nx, ny, nz);
        std::size_t hx = nx / 2, hy = ny / 2, hz = nz / 2;
        for (std::size_t z = 0; z < nz; ++z)
            for (std::size_t y = 0; y < ny; ++y)
                for (std::size_t x = 0; x < nx; ++x)
                    out((x + hx) % nx, (y + hy) % ny, (z + hz) % nz) = input(x, y, z);
        return out;
    }

} // namespace calx

#endif // CALX_FFT_3D_HPP